FastAPIを使用してAPIルーターのテストコードを書く方法について説明します。
特に、テスト用にオンメモリのSQLiteテーブルを使用する際のベクターカラムの扱いについて解説します。
テストコードを作成に必要なパッケージ
まず、FastAPIアプリの依存関係がすでにインストールされていることを前提として、以下のパッケージが必要です。
pytest = "^8.2.2"
aiosqlite = "^0.20.0"
httpx = "^0.27.0"
pytest-asyncio = "^0.23.7"
aiohttp = "^3.9.5"
pytest-aiohttp = "^1.0.5"
テスト対象
以下は、テストしたいルーターの設定を説明します。
データモデルの例
例えば、次のようなデータベースモデルがあるとします。ここで、ベクターカラムが存在していることに注目してください。
class AnkiNoteModel(Base):
"""モデルに基づいてテーブルを作成します。"""
__tablename__ = "anki_cards"
id = Column(Integer, primary_key=True)
deck_name = Column(String, index=True)
front = Column(String, nullable=False)
back = Column(String)
sentence = Column(String)
translated_sentence = Column(String)
audio = Column(String)
vector = Column(Vector(dim=1536))
CRUD関数の例
以下は、指定された単語(データモデルのfront
列)に基づいてカードを取得するCRUD関数の例です。
async def get_card(session: AsyncSession, front: str) -> AnkiNoteModel:
result = await session.execute(
select(AnkiNoteModel).filter(AnkiNoteModel.front == front)
)
return result.scalars().first()
ルーターの例
次に、データベースからアイテムを選択するCRUD関数get_cardを呼び出す簡単なルーターを示します。
# ルーターのレスポンススキーマ
class KoVocResponseSchema(KoVocSchema):
id: int
front: str
back: Optional[str] = None
vector: Optional[List[float]] = None
audio: Optional[str] = None
@router.get("/get_korean_card", response_model=KoVocResponseSchema)
async def get_korean_card(front: str, session: AsyncSession = Depends(get_db)):
card = await get_card(session=session, front=front)
if card:
return card
raise HTTPException(status_code=404, detail="card not found")
テスト関数の作成
テストコードの全文はこちらに:
https://github.com/phlin0424/voc_rec/blob/main/tests/conftest.py
以下は、コードの注目ポイントについて解説します。
テストデータの準備
テスト用データをJSON形式で準備しておきます。例えば、次のようなデータが一つあります。
{
"front": "아버지",
"back": "父親",
"vector": [
0.04205682873725891,
0.017075318843126297,
-0.07281909883022308,
...
],
"audio": "naver-798c54f2-0dca7b9f-a932f140-a2bff2b8-45a4d9c4.mp3",
"id": 73
}
テスト用SQLiteモデルの作成
SQLiteはベクターカラムをサポートしていないため、これを模倣するデータモデルを作成する必要があります。
class AnkiNoteTestModel(SQLiteBase):
"""pgvectorモデルを模倣するSQLiteテーブルモデル"""
__tablename__ = "anki_cards"
id = Column(Integer, primary_key=True)
deck_name = Column(String, index=True)
front = Column(String, nullable=False)
back = Column(String)
sentence = Column(String)
translated_sentence = Column(String)
audio = Column(String)
vector = Column(Text, nullable=False) # PostgreSQLのベクターカラムを模倣
def get_vector(self):
# strを実際のベクターに変換
return json.loads(self.vector)
def set_vector(self, vector):
# ベクターをstrに変換
self.vector = json.dumps(vector)
get_vector()
とset_vector
を注目してください。
テスト用セッションとクライアントの作成
次に、上記のデータモデルを使用してDBを初期化し、セッションとクライアントを作成します。
@pytest_asyncio.fixture(scope="function")
async def async_session(async_db_url) -> AsyncSession:
"""テスト用のSQLite非同期セッションを作成"""
async_engine = create_async_engine(async_db_url, echo=True)
async_session = sessionmaker(
autocommit=False, autoflush=False, bind=async_engine, class_=AsyncSession
)
async with async_engine.begin() as conn:
await conn.run_sync(SQLiteBase.metadata.drop_all)
await conn.run_sync(SQLiteBase.metadata.create_all)
yield async_session
async with async_engine.begin() as conn:
await conn.run_sync(SQLiteBase.metadata.drop_all)
@pytest_asyncio.fixture(scope="function")
async def client(async_session):
"""FastAPIアプリをオーバーライドして非同期クライアントを作成"""
async def _override_get_db():
"""DI override用の関数"""
async with async_session() as session:
yield session
app.dependency_overrides[get_db] = _override_get_db
async with httpx.AsyncClient(app=app, base_url="http://testserver") as client:
yield client
テストデータの挿入
次に、テストデータをDBに挿入します。
@pytest_asyncio.fixture
async def data_in_db(test_data: list[dict], async_session):
"""テストDBにテストデータを挿入"""
async with async_session() as session:
async with session.begin():
for data in test_data:
anki_note = AnkiNoteTestModel(**data)
session.add(anki_note)
await session.commit()
yield test_data
テスト関数の作成
最後に、以下のようにテスト関数を作成します。
@pytest.mark.asyncio
async def test_get_korean_card(client, data_in_db):
response = await client.get("/get_korean_card", params={"front": "아버지"})
assert response.status_code == 200
json_response = response.json()
assert json_response["front"] == "아버지"
assert json_response["back"] == "父親"
assert isinstance(json_response["vector"], list)
# 他のアassertも追加可能です
最後に
FastAPIを使用してAPIルーターのテストコードを書く際に、SQLiteとベクターカラムを扱う方法について説明しました。特に、ベクターカラムのテストをサポートするためのモデル作成やデータベース設定の工夫について紹介しました。
参考記事