Setting up PostgreSQL for running integration tests

When it comes to testing, achieving performance and reliability is crucial. In this article, I'll explain how to set up PostgreSQL for tests and discuss some common pitfalls to avoid.

Isolation as the primary goal

Before we dive into the details, let's define our goals:

• Isolation – We want to ensure that each test runs in isolation. At the very least, this means that each test should have its own database. This ensures that tests don't interfere with each other and that you can run tests in parallel without any issues.
• performance – We want to ensure that setting up PostgreSQL for tests is fast. A slow solution is going to be cost prohibitive for running tests in CI/CD pipelines. The solution that we come up with must allow us to execute tests without introducing too much overhead.

The rest of this article will focus on what we have tried, what worked, and what didn't work.

What didn't work

Using transactions

The first approach we tried was to use transactions. We would start a transaction at the beginning of each test and roll it back at the end.

The basic idea is illustrated in the following example:

test('calculates total basket value', async () => {
  await pool.transaction(async (tx) => {
    await tx.query(sql.unsafe`
      INSERT INTO basket (product_id, quantity)
      VALUES (1, 2)
    `);
 
    const total = await getBasketTotal(tx);
 
    expect(total).toBe(20);
  });
});

The transaction approach works well for simple cases (e.g., testing a single function), but it quickly becomes a problem when dealing with tests that test integration between multiple components. Due to connection pooling, nested transactions, and other factors, the necessary work to make the transaction approach work would have meant that we are not replicating the real-world behavior of our application, i.e. it would not provide the confidence we need.

Using SQLite

Another approach we tried was to use SQLite. SQLite is an in-memory database that is fast and easy to set up.

Similar to the transaction approach, SQLite works well for simple cases. However, it quickly becomes a problem when dealing with code paths that use PostgreSQL-specific features. In our case, due to the use of various PostgreSQL extensions, PL/pgSQL functions, and other PostgreSQL-specific features, we couldn't use SQLite for our tests.

Using pg_tmp

Another approach we tried was to use pg_tmp. pg_tmp is a tool that creates a temporary PostgreSQL instance for each test.

In theory, pg_tmp is a good solution. It allows a complete isolation of tests. In practice, is a lot slower than we could tolerate. With pg_tmp, it takes a few seconds to start and populate the database, and this overhead quickly adds up when running thousands of tests.

Let's say you have 1000 tests, and each test takes 1 second to run. If you add 2 seconds of overhead for creating a new database, you are looking at an additional 2000 seconds (33 minutes) of overhead.

What worked

After trying various approaches, we settled on combining two approaches: template databases and mounting a memory disk. This approach allowed us to isolate each test at a database level without introducing too much overhead or complexity.

Template Databases

A template database is a database that serves as a template for creating new databases. When you create a new database from a template database, the new database has the same schema as the template database. The steps to create a new database from a template database are as follows:

1. Create a template database (ALTER DATABASE <database_name> is_template=true;)
2. Create a new database from the template database (CREATE DATABASE <new_database_name> TEMPLATE <template_database_name>;)

The key advantage of using template databases is that you do not need to mess with managing multiple PostgreSQL instances. You can create copy databases and have each test run in isolation.

However, on its own, template databases are not fast enough for our use case. The time it takes to create a new database from a template database is still too high for running thousands of tests:

postgres=# CREATE DATABASE foo TEMPLATE contra;
CREATE DATABASE
Time: 1999.758 ms (00:02.000)

This is where the memory mounting comes in.

Mounting a memory disk

The final piece of the puzzle is mounting a memory disk. By mounting a memory disk, and creating the template database on the memory disk, we can significantly reduce the overhead of creating a new database.

I will talk about how to mount a memory disk in the next section, but first, let's see how much of a difference it makes.

postgres=# CREATE DATABASE bar TEMPLATE contra;
CREATE DATABASE
Time: 87.168 ms

This is a significant improvement and makes the approach viable for our use case.

Needless to say, this approach is not without its drawbacks. The data is stored in memory, which means that it is not persistent. If the database crashes or the server restarts, the data is lost. However, for running tests, this is not a problem. The data is recreated from the template database each time a new database is created.

Using Docker container with a memory disk

The approach we settled on was to use a Docker container with a memory disk. Here is how you can set it up:

$ docker run \
  -p 5435:5432 \
  --tmpfs /var/lib/pg/data \
  -e PGDATA=/var/lib/pg/data \
  -e POSTGRES_PASSWORD=postgres \
  --name contra-database \
  --rm \
  postgres:14

In the above command, we are creating a Docker container with a memory disk mounted at /var/lib/pg/data. We are also setting the PGDATA environment variable to /var/lib/pg/data to ensure that PostgreSQL uses the memory disk for data storage. The end result is that the underlying data is stored in memory, which significantly reduces the overhead of creating a new database.

Managing test databases

The basic idea is to create a template database before running the tests and then create a new database from the template database for each test. Here is a simplified version of how you can manage test databases:

import {
  createPool,
  sql,
  stringifyDsn,
} from 'slonik';
 
type TestDatabase = {
  destroy: () => Promise<void>;
  getConnectionUri: () => string;
  name: () => string;
};
 
const createTestDatabasePooler = async (connectionUrl: string) => {
  const pool = await createPool(connectionUrl, {
    connectionTimeout: 5_000,
    // This ensures that we don't attempt to create multiple databases in parallel.
    maximumPoolSize: 1,
  });
 
  const createTestDatabase = async (
    templateName: string,
  ): Promise<TestDatabase> => {
    const database = 'test_' + uid();
 
    await pool.query(sql.typeAlias('void')`
      CREATE DATABASE ${sql.identifier([database])}
      TEMPLATE ${sql.identifier([templateName])}
    `);
 
    return {
      destroy: async () => {
        await pool.query(sql.typeAlias('void')`
          DROP DATABASE ${sql.identifier([database])}
        `);
      },
      getConnectionUri: () => {
        return stringifyDsn({
          ...parseDsn(connectionUrl),
          databaseName: database,
          password: 'unsafe_password',
          username: 'contra_api',
        });
      },
      name: () => {
        return database;
      },
    };
  };
 
  return () => {
    return createTestDatabase('contra_template');
  };
};
 
const getTestDatabase = await createTestDatabasePooler();

At this point, you can use getTestDatabase to create a new database for each test. The destroy method can be used to clean up the database after the test has run.

Conclusion

This setup allows us to run thousands of tests in parallel across multiple shards without any issues. The overhead of creating a new database is minimal, and the isolation is at a database level. We are happy with the performance and reliability that this setup provides.