🐳

Vitest + Testcontainers でDBを利用したテストを並列で実行する

に公開

Vitestではデフォルトでテストファイル毎に並列で実行されますが、DBも利用したテストになると、それぞれのテストで同じDBを参照してしまうことでテストが失敗するといったことが発生しかねません。
やもえず並列実行を抑止(fileParallelism: false)して実行していたのですが、テストが多くなるにつれテスト時間の短縮になる並列実行ができないことがネックになってきそうです。

そこで、Testcontainersを利用して並列数分DBコンテナを立ち上げて、それぞれ別々のDBに繋いでテストする方法を試してみることにしました。
Testcontainersは、Dockerコンテナを使ってテスト用の依存サービス(データベースやメッセージキューなど)を起動して利用できるようにするライブラリです。

対象環境

  • パッケージマネージャ: pnpm
  • DB: PostgreSQL
  • ORM: Prisma

実現方法

最終的に下記の方法で実現できました。

  1. VitestのglobalSetupsetupにて、テスト全体での開始時にWorkerの数分Testcontainersを起動する
    • 起動したTestcontainersのDB接続情報を、通番を振った形で環境変数(DATABASE_URL_TEST_1DATABASE_URL_TEST_2...)として保持しておく
    • テーブルの作成まで済ませておく(Prismaのnpx prisma migrate deploy実行)
  2. VitestのsetupFilesで、各テストファイル実行前にDB接続情報を切り替え
    • WorkerのIDを元に、起動時に用意していたDATABASE_URL_TEST_1DATABASE_URL_TEST_2...といった環境変数を参照し、DATABASE_URLに設定しなおす形
  3. 各テストの開始前にデータをリセット
    • beforeEachで、消したくないテーブル以外をTRUCATE TABLE ... RESTART IDENTITY CASCADEでレコード消す&シーケンスリセットする形
  4. VitestのglobalSetupteardownにて、テスト全体での終了時にTestcontainersを停止

コード全体は下記のプロジェクトをご確認ください。DevContianerで動作するようにしています。

実施内容

以降で実際に実施した内容を記載していきます。

(1) @testcontainers/postgresql のインストール

DBがPostgreSQLなので、TestcontainersのPostgreSQL Moduleを利用します。

pnpm add -D @testcontainers/postgresql

なお、コンテナとして起動するので、Dockerが動作している必要があります。

DevContainerならば、docker-in-docker を使うことで簡単にDevContainer上でもDockerが使えて便利です。
下記のように.devcontainer/devcontainer.jsonに追加するだけとなります。

  "features": {
    "ghcr.io/devcontainers/features/docker-in-docker:2": {}
  }

(2) テスト全体の開始/終了時にあわせて、Testcontainerを起動/停止

VitestのglobalSetupで指定したファイルにて、setupteardownでテスト全体の開始、終了時の処理が書けます。

setupにてWorker分のTestcontainerを起動したうえで

  1. Prismaで利用するDB接続文字列を生成
  2. DBのマイグレーション(npx prisma migrate deploy)
  3. DB接続文字列を各Worker用に通番を振った形で環境変数(DATABASE_URL_TEST_1DATABASE_URL_TEST_2...)として退避

といったことを実施するようにしました。
テスト全体の開始時にだけ起動するようにしているのは、Testcontainerの起動/停止の回数を最低限にするためです。
(テストファイルや各テスト毎に起動/停止となると、その分時間がかかることになる)

また、テスト時間が少しでも短縮となるように、PostgreSQLのデータファイルはtmpfs(メモリ上)に配置するよう設定しました。

let containers: StartedPostgreSqlContainer[] = [];

export const setup = async (project: TestProject) => {
  // Worker数分Testcontainerを立ち上げる
  const promises: Promise<StartedPostgreSqlContainer>[] = [];

  console.log("Testcontainer起動開始");
  for (let i = 1; i <= project.config.maxWorkers; i++) {
    // 1から振っていく(process.env.VITEST_POOL_ID と一致するように)
    promises.push(setupTestDatabaseContainer(i));
  }

  containers = await Promise.all(promises);
  console.log("Testcontainer起動完了");
};

const setupTestDatabaseContainer = async (workerId: number) => {
  const container = await new PostgreSqlContainer("postgres:latest")
    // データファイルをtmpfs(メモリ上)に保存することで高速化
    .withTmpFs({ "/var/lib/postgresql/data": "rw" })
    .start();

  // DATABASE_URLを作成
  const databaseUrl = `postgresql://${container.getUsername()}:${container.getPassword()}@${container.getHost()}:${container.getMappedPort(
    5432
  )}/${container.getDatabase()}`;

  execSync(`DATABASE_URL=${databaseUrl} npx prisma migrate deploy`);
  console.log(`(workerId:${workerId}) DATABASE_URL: ${databaseUrl}`);

  process.env[`DATABASE_URL_TEST_${workerId}`] = databaseUrl;

  return container;
};

teardownでは、起動済みのTestcontainerを停止していきます。

export const teardown = async () => {
  console.log("Testcontainer停止開始");
  // Testcontainerを停止する
  await Promise.all(containers.map((container) => container.stop()));
  console.log("Testcontainer停止完了");
};

(3) 各テストファイル実行時にDB接続情報を切り替え

VitestのsetupFilesで指定したファイルで、各テストファイル実行前にDB接続情報を切り替えます。
Testcontainer起動時に各Worker毎に環境変数として退避済み(DATABASE_URL_TEST_1DATABASE_URL_TEST_2...)なので、それをDATABASE_URLに設定しなおす形です。(PrismaではDATABASE_URLを使ってDB接続が行われる)

Workerを特定する方法として、環境変数のVITEST_POOL_IDが利用できるので、それを元に設定します。

// 各Worker用のDB接続情報へ切り替え
process.env.DATABASE_URL =
  process.env[`DATABASE_URL_TEST_${process.env.VITEST_POOL_ID}`];

(4) 各テストの開始前にDBをリセット

起動したTestcontainerを使いまわすことになるので、各テストで登録したデータが残ったままになります。
これを各テスト開始前にリセットするため、beforeEachTRUCATE TABLE テーブル1, テーブル2, ... RESTART IDENTITY CASCADEでレコード消す&シーケンスをリセットします。

_prisma_migrationsは、Prismaの管理テーブルとなり、消す必要が無いので対象外にしています。

// テスト毎にDBリセット
const prisma = new PrismaClient();
await prisma.$connect();

const resetDb = async () => {
  const tablenames = await prisma.$queryRaw<
    Array<{ tablename: string }>
  >`SELECT tablename FROM pg_tables WHERE schemaname='public'`;

  const tables = tablenames
    .map(({ tablename }) => tablename)
    // _prisma_migrationsはマイグレーションの履歴テーブルなので除外
    // 他にも除外したいテーブルがあればここに追加
    .filter((name) => name !== "_prisma_migrations")
    .map((name) => `"public"."${name}"`)
    .join(", ");

  try {
    await prisma.$executeRawUnsafe(
      `TRUNCATE TABLE ${tables} RESTART IDENTITY CASCADE;`
    );
  } catch (error) {
    console.log({ error });
  }
};

beforeEach(async () => {
  await resetDb();
});

(5) Testcontainerの起動にかかる時間を考慮してhookTimeoutを延ばす

フックのタイムアウトがデフォルトだと10秒になっており、Testcontainerの起動が10秒で終わらずタイムアウトしかねないので、時間を延ばしておきます。
hookTimeoutで変更できます。

import { defineConfig } from "vitest/config";

export default defineConfig({
  test: {
    globals: true,
    globalSetup: ["./test/global-setup.ts"],
    setupFiles: ["./test/setup.ts"],
    watch: false,
    hookTimeout: 600000, // Testcontainerの起動に初回は時間がかかる可能性があるため
    minWorkers: 4,
    maxWorkers: 4,
  },
});

終わりに

この方法を使うことで、DBを使ったテストもお互い競合することなく並列で実行できるようになるので、テスト時間の短縮の手段として使える目途が立ちました。
ただ、Testcontainerの起動にはそれなりに時間がかかりますし、リソースも消費することになるので、バランスを見ながら(並列数も考えながら)試してみることになりそうです。

Discussion