【Python・FastAPI】PyTest + FactoryBoy + SQLAlchemyでデータベーステストをする方法

概要

FastAPIの開発において、より本格的に、実際の環境に近い状態でテストを行いたかったので、テスト用のデータベースを作成し、そこにfactory_boyで作成したデータを挿入してテストを行うことにしました。

実装環境

• python: 3.11.0

• PostgreSQLを使用

• 新しいsqlalchemy2.0の記法を極力使用

• インフラはdocker composeを使用（docker composeを使っていなくても参考になると思います）

• プロジェクトディレクトリ構成:

.
├── README.md
├── __pycache__
├── alembic.ini
├── config
├── cruds
├── main.py
├── migrations
├── models
├── public
├── requirements.txt
├── routers
├── schemas
└── tests

• tests配下のディレクトリ構成

├── __init__.py        
├── conftest.py         # こちらにfixtureを記述
├── factories
│   └── user.py         # FactoryBoyのモデル定義ファイル
└── requests
    ├── __init__.py     
    └── test_user.py    # APIリクエストのテスト

実装

1. conftest.pyにfixtureを定義

tests/conftest.py

import pytest
from fastapi.testclient import TestClient
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy import create_engine
import os
# 本来のデータベースの設定ファイルからBase, get_dbをインポート
from config.database import Base, get_db
from main import app

# テストデータベースの情報
TEST_POSTGRES_USER = *****
TEST_POSTGRES_PASSWORD = *****
TEST_POSTGRES_DB = *****
TEST_POSTGRES_PORT = ****

# SQLAlchemyを用いたデータベース接続のセッティング
SQLALCHEMY_DATABASE_URL = f"postgresql://{TEST_POSTGRES_USER}:{TEST_POSTGRES_PASSWORD}@test_db:{TEST_POSTGRES_PORT}/{TEST_POSTGRES_DB}"
engine = create_engine(SQLALCHEMY_DATABASE_URL)
TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
ScopedSession = scoped_session(TestingSessionLocal)

# テスト時にget_dbをオーバーライドする関数
def override_get_db():
    try:
        db = TestingSessionLocal()
        yield db
    finally:
        db.close()

# テストの最初だけ実行するfixtureを定義(*scope="session"より)
@pytest.fixture(autouse=True, scope="session")
def client():
    # テストDBにテーブルなどを作成
    Base.metadata.create_all(bind=engine)
    # get_dbをテスト時のみオーバーライド
    app.dependency_overrides[get_db] = override_get_db
    
    with TestClient(app) as c:
        yield c

    # テストDBのリセット
    Base.metadata.drop_all(bind=engine)


# 関数単位で実行されるfixture
@pytest.fixture(scope="function")
def session():

    # ここで先ほど定義したScopedSessionをyield。他のテストファイルで使い回す。
    yield ScopedSession

    ScopedSession.close()

2. model、factoryの実装
今回はデータベースモデルのサンプルとしてUserを定義します。SQLAlchemy2.0の記法を使用していますが、型ヒントがつく以外同じなので従来の記法でも同じです。

models/user.py

from sqlalchemy import BigInteger, Boolean, String
from sqlalchemy.orm import Mapped, mapped_column

from config.database import Base


class User(Base):
    __tablename__ = "users"
    
    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, index=True)
    is_deleted: Mapped[bool] = mapped_column(Boolean, default=False)
    name: Mapped[str] = mapped_column(String(100))
    hashed_password: Mapped[str] = mapped_column(String(100))

tests/factories/user.py

from factory.alchemy import SQLAlchemyModelFactory
from factory import Faker
from tests.conftest import ScopedSession

from models.user import User

class UserFactory(SQLAlchemyModelFactory):
    class Meta:
        # ここでモデルを指定
        model = User
        # conftest.pyで定義したScopedSessionをここで定義する
        sqlalchemy_session = ScopedSession
        sqlalchemy_session_persistence = "commit"

    # Fakerを用いてカラムをそれぞれ定義
    id = Faker("random_int", min=1)
    is_deleted = False
    is_owner = False
    is_admin = False
    name = Faker("name")
    hashed_password = Faker("password")

3. テストコードの作成
tests/requests/test_user.pyにテストを書いていきます。今回はFactoryBoyで作成したUserのインスタンスがデータベースに保存されているか確認するテストを書きます。今回ルーティング、CRUDなどは本題ではないため本来スケーリングの観点からファイルを分けて書くべき処理を一つのファイルに書かせていただきます。

/usersにアクセスするとUserが一覧取得できるエンドポイントを作成します。

main.py

from fastapi import FastAPI

app = FastAPI()

@app.get("/users")
def get_users(db: Session = Depends(get_db)):
    # Userを全件返すだけのAPI
    return db.query(User).all()

tests/test_user.py

import pytest
from fastapi.testclient import TestClient
from main import app
from tests.factories.user import UserFactory

client = TestClient(app)

# 引数にsessionを指定することでconftest.pyで定義したsessionのfixtureを使う
def test_get_admin_users(session):

    # Factoryのインスタンスを作成。
    # ここで定義した通りに適当なデータがuserに代入されています。
    # その後このインスタンスは自動でテストDBに保存されます。
    user = UserFactory()
    
    # APIリクエストの送信
    response = client.get("/users")
    assert len(response.json()) == 1
    assert response.json() == [{"id": user.id, "name": user.name, "is_owner": user.is_owner, "is_admin": user.is_admin}]
    assert response.status_code == 200

4. テストの実行

$ pytest

最後に

fixture内で必死に定義する必要があったり、テストコードが長くなる原因となっていたテストデータの作成をオブジェクト指向的に管理できてコードも短くできるので、とてもいいアプローチだと思っています🙂‍↕️読んでいただきありがとうございました。