How to mock node-postgres in vitest

245 views Asked by At

I am trying to mock node-postgres in vitest with esm. But I am running into different problems. When I am using the Pool class in my function:

import {Pool} from 'pg';
...
const pool = new Pool({connectionString});
const client = await pool.connect();
await client.query('BEGIN');
...

When I am importing Pool as above with import {Pool} from 'pg' i can mock the pool with:

vi.mock('pg', () => {
    const Pool = vi.fn();
    const Client = vi.fn();
    return {Pool, Client};
});
function queryMockImplementation(migrationTableExists: boolean, migratedFiles?: string[]) {
    return (queryString: string) => {
        if (queryString === `SELECT EXISTS (SELECT FROM pg_tables WHERE schemaname = 'administration' AND tablename = 'migrations')`) {
            return Promise.resolve({rows: [{exists: migrationTableExists}]});
        }
        return Promise.resolve({rows: []});
    }
}
function connectMockImplementation(clientQueries: string[], throwOnSqlQuery?: string) {
    return async() => {
        return {
            query: vi.fn().mockImplementation((queryString: string) => {
                if (throwOnSqlQuery === queryString) {
                    throw new Error('Error thrown by mock for client query.');
                }
                clientQueries.push(queryString);
            }),
            release: vi.fn(),
        };
    };
}

Then in the test I can make the following:

it('Should mock', () => async {
            const {Pool} = await import('pg');
            const clientQueries: string[] = [];
            Pool.prototype.query = vi.fn().mockImplementation(queryMockImplementation(false));
            Pool.prototype.connect = vi.fn().mockImplementation(connectMockImplementation(clientQueries));
            const pool = new Pool();
}

But this is not the intended way. It works but I have to inject the array clientQueries into the mock implementation for the client to observe the strings called with client.query. I could not find a way to get the client directly mocked. Also this way of importing Pool creates problems with esm in packages consuming this library. So when I try to access Pool with:

import pg from 'pg';
const {Pool} = pg;

I get the vitest error:

Error: [vitest] No "default" export is defined on the "pg" mock. Did you forget to return it from "vi.mock"?
If you need to partially mock a module, you can use "vi.importActual" inside:

vi.mock("pg", async () => {
  const actual = await vi.importActual("pg")
  return {
    ...actual,
    // your mocked methods
  },
})

So I am adjusting the mock:

vi.mock('pg', async () => {
    const actual = <Record<string, any>>await vi.importActual("pg")
    const Pool = vi.fn();
    const Client = vi.fn();
    return {...actual, Pool, Client};
});

But now I have two other problems. I have to provide an actual database and connection string in the test and the client query mock is not working any more.

1

There are 1 answers

0
Florat On BEST ANSWER

I got the issue fixed. Example repository at vitest-allowSyntheticDefaultImports

import {describe, vi, it, expect} from 'vitest';
import {connect} from './index';
import pg from 'pg';
const {Pool} = pg;

vi.mock('pg', async () => {
    const actual = <Record<string, unknown>>await vi.importActual('pg');

    return {
        ...actual,
        default: {
            Pool: vi.fn(),
        }
    };
});

describe('connect', () => {
    it('should call Pool', () => {
        connect('postgres://...');
        expect(Pool).toBeCalledTimes(1);
    });
});