やること
本記事で使用したコードはGitHubにあります。Poetryが入っていればすぐに動かせます。
1. FastAPIでsqlite3を使用した簡単なCRUDアプリケーションの実装
CRUDアプリケーション
userアカウントに関するCRUDアプリを作成します。テスト手法の紹介がメインなので以下の簡単な機能だけ実装します。
- Create: ユーザー登録
準備
FastAPIに加えて以下のパッケージのpip installが必要です。
- sqlalchemy
- sqlalchemy-utils
- async-exit-stack (Python 3.7では不要)
- async-generator (Python 3.7では不要)
- requests
- pytest
ディレクトリ構成
以下の様なディレクトリ構成で必要なものを実装します。
├── tests
│ ├── __init__.py
│ ├── conftest.py # pytestのfixture定義
│ └── test_user.py # APIのテスト
└── users
├── __init__.py
├── crud.py # クエリ発行用の関数定義
├── database.py # データベース設定
├── main.py # API定義
├── models.py # table定義
└── schemas.py # APIのI/O定義
データベースの設定 (database.py)
Database URLで接続するデータベースの指定をします。基本的に環境変数で宣言すべきですが、簡単のためにベタ書きします。
主要なDatabaseのURLの記法はこちらにまとまっています。
connect_args={"check_same_thread": False}の部分はsqlite3用の設定なので、他のDatabaseを使用する場合は削除して下さい。
SessionLocal変数に入っているsessionmaker instanceは,
callするとsession instanceを生成します。これは、DBとのコネクションを管理するために使います。また、SQLクエリの発行にもsessionを使います。
import os
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
SQLALCHEMY_DATABASE_URL = os.environ.get('DATABASE_URL', 'sqlite:///./test.db')
engine = create_engine(
SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
Table定義 (models.py)
databaseの設定時に定義したBaseを継承してTable定義を行います。この様に定義すると、Base経由でtableの作成やORMマッパーの使用が簡潔に行えるようになります。
from sqlalchemy import Boolean, Column, Integer, String
from .database import Base
class User(Base):
__tablename__ = "users"
id = Column(Integer, primary_key=True, index=True)
email = Column(String, unique=True, index=True)
hashed_password = Column(String)
is_active = Column(Boolean, default=True)
APIのI/O定義 (schemas.py)
APIのI/Oを定義します。ここでは、emailとpasswordを入力 -> 作成されたuserのid, email, アクティブユーザーかどうかを返すことにします。schemaを決めておくだけでSerializeやDeserializeを勝手に行ってくれます。
from pydantic import BaseModel
class UserBase(BaseModel):
"""Base User scheme"""
email: str
class UserCreate(UserBase):
"""Input"""
password: str
class User(UserBase):
"""Output"""
id: int
is_active: bool
class Config:
orm_mode = True
クエリ発行用の関数定義 (crud.py)
sqlalchemyではsessionを使ってSQLクエリを発行します。この処理は問題が起きやすいので切り出して単体テストを行いやすくしておきます。極力、ロジックは入れずにsessionを受け取ってクエリ発行だけ行うようにします。
from sqlalchemy.orm import Session
from hashlib import md5 as hash_func
from . import models
from . import schemas
def get_user_by_email_query(db: Session, email: str):
"""get user by email"""
return db.query(models.User).filter(models.User.email == email).first()
def create_user_query(db: Session, user: schemas.UserCreate):
"""create user by email and password"""
hashed_password = hash_func(user.password.encode()).hexdigest()
db_user = models.User(email=user.email, hashed_password=hashed_password)
db.add(db_user)
db.commit()
db.refresh(db_user)
return db_user
API定義 (main.py)
CRUD APIを定義します。
注意する必要があるのは、sessionの渡し方です。引数にdependsとして関数やクラスを宣言すると、それをcallした結果(関数ならreturn、classならinstance)が引数に渡されます。これを利用して、request毎にSessionLocalを使ってsessionを生成し、databaseとのコネクションを確保しています。そして、そのsessionを用いてクエリを発行しています。
from fastapi import Depends, FastAPI, HTTPException
from sqlalchemy.orm import Session
from . import models, schemas
from .crud import (
get_user_by_email_query,
create_user_query
)
from .database import SessionLocal, engine
# table作成
models.Base.metadata.create_all(bind=engine)
app = FastAPI()
# Dependency
def get_db():
try:
db = SessionLocal() # sessionを生成
yield db
finally:
db.close()
@app.post("/users/", response_model=schemas.User)
def create_user(user: schemas.UserCreate, db: Session = Depends(get_db)):
db_user = get_user_by_email_query(db=db, email=user.email)
if db_user:
raise HTTPException(status_code=400, detail="Email already registered")
return create_user_query(db=db, user=user)
なお、sessionを渡す方法としてmiddlewareを使った実装もありますが、すべてのAPIでDBとのconnectionをつくってしまうので、DBを使わないAPIが多い場合は無駄が多くなるなどの弊害があるため現在は非推奨だそうです。(参考)
2. テストケース毎にDBを作成し、他のテストケースや本番用のDBに影響を与えずにAPIのunittestを行う
本題に入ります。
APIのテスト (本番用のDBを使用する場合)
FastAPIではstarlette.testclient.TestClientで、以下のように簡潔にAPIのテストができます。
from starlette.testclient import TestClient
from users.main import app
client = TestClient(app)
def test_create_user():
response = client.post(
"/users/", json={"email": "foo", "password": "fo"}
)
assert response.status_code == 200
ここで、pytestを実行すると自動テストが行なえます。
$ pytest tests
しかし、APIは本番用のDBに接続してしまうのでテストを実行すると、userを追加してしまいます。2回テストを実行すると、2回目は既に同名のemailが登録されているのでuserの作成に失敗してテストが通らなくなります。
そこで、テスト実行時に一時的にDatabaseを作成し、本番用のDatabaseに影響がでないようにし、かつ、テスト実行毎に毎回クリーンなDatabaseを用意できるようにします。さらに、汎用的に使えるように関数単位でDatabaseを作り直し、テストケース毎に互いに影響がでないようにします。
テスト用のクリーンなDBの作成と削除
クリーンなDBでテストを行うために必要な処理は、以下のようなものです。
- 関数毎に一時的にDatabaseを作成
- そのDatabaseとのsessionをつくることができるSessionmakerインスタンスをテスト関数に渡す
- 各テストケースの終了時にDatabaseを削除
これらができれば、テストケース毎にクリーンなDatabaseを利用することができて、終了時には何も痕が残りません。
この処理はテストケースに依らず必要になってくるので、このような処理を行うfixtureをconftest.pyの中で定義することにします。
実装は以下のようになります。
import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy_utils import database_exists, drop_database
from .database import Base
@pytest.fixture(scope="function")
def SessionLocal():
# settings of test database
TEST_SQLALCHEMY_DATABASE_URL = "sqlite:///./test_temp.db"
engine = create_engine(TEST_SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False})
assert not database_exists(TEST_SQLALCHEMY_DATABASE_URL), "Test database already exists. Aborting tests."
# Create test database and tables
Base.metadata.create_all(engine)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
# Run the tests
yield SessionLocal
# Drop the test database
drop_database(TEST_SQLALCHEMY_DATABASE_URL)
テスト時にAPIが接続するDBを変更させる
fixtureのおかげで引数にSessionLocalを宣言すると、その関数の実行時に、クリーンなDatabaseが作られるようになりました。後は、APIの接続するDBを強制的にテスト用のものに変更する必要があります。影響を減らすためにテストコードだけで完結させたいです。
FastAPIでは、APIの引数に宣言されたFastAPI.Dependsはapp.dependency_overridesで強制的に上書きできます。
なので、main.get_dbを上書きしてテスト用のsessionmaker instanceを使うように書き換えれば接続先を変更できます。
そこで、以下のようなdecoratorを定義します。
from users.main import app, get_db
def temp_db(f):
def func(SessionLocal, *args, **kwargs):
# テスト用のDBに接続するためのsessionmaker instanse
# (SessionLocal) をfixtureから受け取る
def override_get_db():
try:
db = SessionLocal()
yield db
finally:
db.close()
# fixtureから受け取るSessionLocalを使うようにget_dbを強制的に変更
app.dependency_overrides[get_db] = override_get_db
# Run tests
f(*args, **kwargs)
# get_dbを元に戻す
app.dependency_overrides[get_db] = get_db
return func
テストコードを修正し、さきほど定義したdecoratorを使うだけで、関数単位で、テスト実行時にテスト用の一時的なDatabaseを作成し、そのDatabaseを使ってテストを行えるようになります。なお、それぞれ独立なDatabaseを使うことになるので他のテストケースに影響を及ぼしません。
from starlette.testclient import TestClient
from users.main import app
client = TestClient(app)
@temp_db
def test_create_user():
response = client.post(
"/users/", json={"email": "foo", "password": "fo"}
)
assert response.status_code == 200
おわりに
FastAPIでテスト用のDBを作成してpytestでAPIのUnittestを行う方法をまとめました。
テストケース毎にDBを作り直すのは処理速度が遅くなりそうなので、rollbackする方法を試行錯誤していましたが断念してしまいました。(追記: この方法はこちらで紹介されています)
今回の方法でもテストケースが少ない、または、テストする際のデータ量が多くない場合は有効だと思うのでこの記事が何かの参考になれば幸いです。