SQLModelでDBモックを生成する

はじめに

個人でサクッと作ったPythonプログラムに、SQLModel を使っています。 テストコードを書く際にモック化したことをメモしていきます。

前提

DBは PostgreSQL を利用する想定です。 また、バージョン等は以下の通りになります。

• SQLModel 0.0.16
• PostgreSQL 16.2

DBに接続する場合

さて、通常は以下のように create_engine メソッドを使ってDBとのコネクションを確立し、SQLを実行します。

import typing
import sqlmodel


# テーブルモデル
class Todo(sqlmodel.SQLModel, table=True):
    id: typing.Optional[int] = sqlmodel.Field(default=None, primary_key=True)
    title: str

# コネクション確立
engine = sqlmodel.create_engine(
    url="postgresql+psycopg2://user:password@host:5432/db_name",
    echo=False
)

# SQL実行
with sqlmodel.Session(engine) as session:
    session.add(Todo(title="テスト１"))
    session.add(Todo(title="テスト２"))

ただ、これだとテストコードを書く際にDBを直接接続してしまい、余分なデータが溜まってしまいます。 その時に役立つのモックです。

モックを利用する場合

モック化する場合、create_mock_engine メソッドを使い、Session クラスに bind します。

import typing
import sqlmodel


# テーブルモデル
class Todo(sqlmodel.SQLModel, table=True):
    id: typing.Optional[int] = sqlmodel.Field(default=None, primary_key=True)
    title: str

# コネクション確立をモック化
engine = sqlmodel.create_mock_engine(
    url="postgresql+psycopg2://",
    executor=None
)

# SQL実行(※モックのためDBにはデータは登録されない)
with sqlmodel.Session() as session:
    session.bind = engine
    session.add(Todo(title="テスト１"))
    session.add(Todo(title="テスト２"))

おわりに

SQLModelがモック用のメソッドを提供しているため、簡単にモックを作成することができました。 この機能を活用することで、テストコードをより気軽に記述できるようになります。