Drizzle, Jestで実DBを使ったテストを並列で動かす

Drizzle ORMというタイプセーフで軽量なORMライブラリがあります。Edge環境での動作がサポートされており、PlanetScaleなどのDBに接続できることから、個人的に好んで利用しています。
https://orm.drizzle.team/
しかし、まだマイナーなライブラリであるためエコシステムが育っておらず、例えばPrismaでいうところのjest-prismaのような実際のDBを使用したテストを容易に行うためのツールが存在しません。
https://github.com/Quramy/jest-prisma

この記事では、Drizzle, Jestで実際のDBに接続し、テストケースごとにDBの状態を初期化しつつ並列でテストを実行する方法を解説します。

Jestの設定

前述したjest-prismaのコードを参考にしています。各テストケースの実行前にトランザクションを張り、実行後にロールバックしてあげる、という考え方です。

以下はJestのsetupFilesAfterEnvで指定するファイルです。
この例ではMySQLを使用しています。importは適宜使用しているDBに置き換えてください。
また、Drizzleクライアントは以下のようにシングルトン的にモジュールからexportしていることを前提としています。

src/db/index.ts

import { drizzle } from 'drizzle-orm/mysql2'
import mysql from 'mysql2/promise'
import * as schema from './schema'

const connection = mysql.createPool(process.env.DATABASE_URL!)
export const db = drizzle(connection, { schema: schema, mode: 'default' })

jest.setup.ts

import { drizzle, MySql2Database } from 'drizzle-orm/mysql2'
import mysql from 'mysql2/promise'
import { sql } from 'drizzle-orm'
import * as schema from '@/db/schema'

jest.mock('@/db', () => ({
  db: undefined,
}))

let connection: mysql.Connection
let db: MySql2Database<typeof schema>
let rollbackTransaction: () => void

beforeAll(async () => {
  connection = await mysql.createConnection(process.env.DATABASE_URL!)
  db = drizzle(connection, { schema: schema, mode: 'default' })
})

afterAll(() => {
  connection.end()
})

beforeEach(
  () =>
    new Promise<void>((resolve) => {
      db.transaction((tx) => {
        const originalDb = require('@/db')
        originalDb.db = tx
        resolve()

        return new Promise((_, reject) => {
          rollbackTransaction = reject
        })
      }).catch(() => {
        // ignore
      })
    })
)

afterEach(() => {
  rollbackTransaction()
})

以下で簡単に解説します。

jest.mock('@/db', () => ({
  db: undefined,
}))

この部分で本来使用しているDrizzleのクライアントをモックしています。

beforeEach(
  () =>
    new Promise<void>((resolve) => {
      db.transaction((tx) => {
        const originalDb = require('@/db')
        originalDb.db = tx
        resolve()

        return new Promise((_, reject) => {
          rollbackTransaction = reject
        })
      }).catch(() => {
        // ignore
      })
    })
)

beforeEach内で、つまり各テストケースの実行前にトランザクションを張り、本来のDrizzleクライアントを置き換えます。
beforeEachのコールバックではPromiseを返すようにし、Drizzleクライアントを置き換えたあとに解決します。
トランザクション内でPromiseをrejectしてあげればロールバックされるので、rollbackTransactionにrejectを割り当てます。

afterEach(() => {
  rollbackTransaction()
})

そしてafterEach内で、つまり各テストケースの実行後にロールバックさせます。

---

あとはこのファイルをjest.configのsetupFilesAfterEnvで指定し、適当にローカルでDBを立ち上げ、テストを実行しましょう。

おわりに

書いてみれば簡単で、またほぼjest-prismaとやっていることは同じですが、Jestのモックの方法に慣れていなかったりして少し手間取りました。
DrizzleはPrismaのようにORM寄りではなく、薄めのクエリビルダーのようなものなので使い勝手は異なりますが、Edge環境で動くのは個人的に大きなメリットです。Next.jsのApp Router（サーバーコンポーネント）から手軽にDBと接続したいときなどに使いやすいと思います。興味を持った方は利用してみてはいかがでしょうか。