DBを使ったテストを超高速化する

2024/02/23に公開

自己紹介

私はSaaS系小規模スタートアップ(エンジニア2人)でエンジニアをしているものです。
元々フロントエンドエンジニアとして入社しましたが、その後前任者が辞めたのをきっかけにAWSを含めた開発の全範囲を担当するようになり1年ちょっとが経ちました。
Zennへの投稿は初めてで、技術記事自体もかなり久々にかきました。
性格が雑なので記事も雑です。すみません。
あーしてこーしてとかあれば指摘してください!

結論

この方法を使うと私の環境ではほぼDBアクセスのオーバーヘッドを感じなくなるレベルになりました。
MySQLをつかっていますが、PostgreSQLでも同様のことができると思います。

はじめに

使用技術: TypeScript, Prisma, MySQL, Docker, Vitest
環境: Ubuntu

SQLのクエリが複雑になる場合がありクエリも含めてテストしたい。ただDBを使用したテストは遅いし、DBの状態に依存するので壊れやすい。
今回はこの問題をクリアした最強()の方法を思いついたので共有します。
もっといい方法があれば教えてください。
以下のポイントを押さえれば、言語やライブラリが違っても応用できます。

  1. tmpfsを使う
  2. 事前に用意した複数のdatabaseを使用
  3. 高速なテーブルリセット

この記事ではすでにテスト、Prisma Schemaが存在し、諸設定が完了している前提で説明します。

docker composeでインスタンスを起動

今回のテスト用DBを専用で作成します
portは3307にしているので注意
ポイントはtmpfsを使用すること

docker-compose.yml
version: 3.9
services:
  mysql-test:
    image: mysql:8.0
    restart: always
    environment:
      MYSQL_ROOT_PASSWORD: root
    ports:
      - "3307:3306"
    tmpfs:
      - /var/lib/mysql
docker compose up -d

これでテスト専用mysqlインスタンスが立ち上がります。

事前に複数のdatabaseを作成

この場合100個作成するので結構時間がかかります。
必要なライブラリは事前にインストールしてください。
ここではPrismaを使用していますが、他のライブラリを使っている場合は一部変えてください。

init_testdb.ts
import mysql, { createConnection } from "mysql2/promise";
import { exec } from "child_process";
import fs from "fs/promises";

// 100個の配列を作成
const DB_INDIES = new Array(100).fill(null).map((_, i) => i);

export const initTestDb = async () => {
  // クライアント初期化
  const client = await mysql.createConnection({
    host: "localhost",
    port: 3307,
    user: "root",
    password: "root",
  });

  const start = Date.now();
  console.log("Initializing test databases...");
  await Promise.all(
    DB_INDIES.map(async (index) => {
      // DB再作成
      await client.query(`DROP DATABASE IF EXISTS test_${index}`);
      await client.query(`CREATE DATABASE test_${index}`);

   const DB_URL = `mysql://root:root@localhost:3307/test_${index}`;
      // スキーマを適用。Prisma以外を使っている場合はここを変更してください。
      const { stderr } = await execAsync(
        `PRISMA_DB_URL="${DB_URL}" npx prisma db push --schema=./prisma/prisma/schema.prisma`,
      );
      if (stderr) {
        throw new Error(stderr);
      }
    }),
  );

  await client.end();

  console.log(`Test databases initialized in ${Date.now() - start}ms`);
};

initTestDb()

ここでtsxを使います。tsxはts-nodeの改良版みたいなもので、簡単にTypesScriptで書いたコードを実行できるため便利です

npx tsx init_testdb.ts

これで同じ構造のdatabaseが100個作成されます。

databaseの排他制御

複数databaseを作成しましたが、普通に使うと使用するdatabaseを重複しないように管理する必要が出てきます。
テストファイルごとにdatabaseを作成してもいいのですが、管理がだるいので少し頑張ります。
今回は、proper-lockfileを使ってファイルロックによる排他制御を行い、同時に同じdatabaseを使えないようにします。
私はvitestを使っていますが、jest等でも大体同じだと思います。

流れとしては

  1. globalSetupで、テスト実行前に各databaseに対応するファイルを生成
  2. setupfilesで各テストファイルの前に利用可能なdatabase名を取得
  3. 各テストファイル実行後に利用したdatabaseをリセット&解放
    といった感じです。
testdb-utils.ts
import { exec } from "child_process";
import fs from "fs/promises";
import mysql, { createConnection } from "mysql2/promise";
import lockFile from "proper-lockfile";
import { promisify } from "util";

const execAsync = promisify(exec);
const DB_SEMAPHORE_DIR = "/tmp/.test_db_semaphore";
const DB_COUNT = 100;
const DB_INDIES = new Array(DB_COUNT).fill(null).map((_, i) => i);
const PORT = 3307;

// ファイルロックにより排他制御
export const initSemaphore = async () => {
  await fs.rm(DB_SEMAPHORE_DIR, { recursive: true, force: true });
  await fs.mkdir(DB_SEMAPHORE_DIR);
  await Promise.all(
    new Array(DB_COUNT)
      .fill(null)
      .map((_, i) => fs.writeFile(`${DB_SEMAPHORE_DIR}/${i}`, "")),
  );
};

//現在利用可能なdatabaseを取得
export const getAvailableDbName = async () => {
  let errorCount = 0;
  // 一応リトライ処理を書いておく
  while (true) {
    const semaphore = await fs.readdir(DB_SEMAPHORE_DIR);
    if (semaphore.length === 0) {
      throw new Error("No available db found");
    }

    const lockStatuses = await Promise.all(
      semaphore.map(async (s) => ({
        name: s.replace(".lock", ""),
        locked: await lockFile.check(
          `${DB_SEMAPHORE_DIR}/${s.replace(".lock", "")}`,
        ),
      })),
    );
    const unlocked = lockStatuses
      .filter((status) => !status.locked)
      .map((status) => status.name);

    try {
      if (unlocked.length === 0) {
        throw new Error("No available db found");
      }

      const target = unlocked.pop();
      await lockFile.lock(`${DB_SEMAPHORE_DIR}/${target}`);
      await resetDatabaseStructure(`test_${target}`);
      return `test_${target}`;
    } catch (e) {
      errorCount++;
      if (errorCount > 10) {
        throw e;
      }

      // ちょっとまつ
      await new Promise((resolve) => setTimeout(resolve, 100));
      continue;
    }
  }
};

// 使用後に解放
export const releaseDb = async () => {
  const db = process.env.DATABASE_URL?.split("/").pop();
  if (db == null) {
    throw new Error("DATABASE_URL is invalid");
  }

  const semaphore = db.split("_").pop();

  if (semaphore == null) {
    throw new Error("DATABASE_URL is invalid");
  }

  const index = Number.parseInt(semaphore, 10);

  await resetDatabaseStructure();
  await lockFile.unlock(`${DB_SEMAPHORE_DIR}/${index}`);
};

// DBを高速でリセットする関数 もっといいのがあったら教えてください。
const resetDatabaseStructure = async (
  databaseUrl: string = process.env.DATABASE_URL ?? "",
) => {
  const client = await createConnection({
    host: "localhost",
    port: PORT,
    user: "root",
    password: "root",
  });

  try {
    const start = Date.now();
    const dbName = databaseUrl?.split("/").pop();
    if (dbName == null) {
      throw new Error("DATABASE_URL is invalid");
    }
    await client.query(`USE ${dbName}`);
    await client.query("SET FOREIGN_KEY_CHECKS = 0"); // 外部キー制約を一時的に無効にする
    const [tables]: any[] = await client.query("SHOW TABLES");
    const tableNames = tables.map((table) => Object.values(table)[0]);

    const createTableQueries = [];

    for (const tableName of tableNames) {
      const [createInfo]: any = await client.query(
        `SHOW CREATE TABLE ${tableName}`,
      );
      // @ts-expect-error
      createTableQueries.push(createInfo[0]["Create Table"]);
    }

    for (const tableName of tableNames) {
      await client.query(`DROP TABLE IF EXISTS ${tableName}`);
    }

    for (const query of createTableQueries) {
      await client.query(query);
    }
  } catch (error) {
    console.error(error);
  } finally {
    // foreign key checksを元に戻す
    await client.query("SET FOREIGN_KEY_CHECKS = 1");
    await client.end();
  }
}

globalSetupファイル内でinitSemaphore()を実行

globalSetup.ts
import { initSemaphore } from "./testdb-utils";

export default async () => {
  await initSemaphore();
};

各テストファイルの前に実行されるsetup.tsでPrisma Clientで使用される環境変数を設定する。
そしてテスト終了後に解放。

setup.ts
import { getAvailableDbName, releaseDb } from "./testdb-utils";

const testDBName = await getAvailableDbName();
process.env.DATABASE_URL = `mysql://root:root@localhost:3307/${testDBName}`;

afterAll(async () => {
  await releaseDb();
});

vitest設定ファイル

vitest.config.ts
import devServer from "@hono/vite-dev-server";
import tsconfigPaths from "vite-tsconfig-paths";
import { defineConfig } from "vitest/config";

export default defineConfig({
    include: ["src/**/*.spec.ts"],
    globals: true,
    alias: {
      "@": "src",
    },
    setupFiles: ["./setup.ts"],
    globalSetup: "./globalSetup.ts",
    typecheck: {
      enabled: false,
    },
});

あとは適当にテストを作って実行してください。

最後に

普段技術記事を書くことはあまりなかったのですが、今回三連休の勢いで書いてます。
性格上雑な記事になってしまいますが、バズることを目標にせず、今後のキャリアのための貯蓄として書き溜めていこうとおもいます。
間違っていたら教えてください。

Discussion