0
0

FastAPIでAPIルーターのテストコードを書く方法:SQLiteとベクターカラムの扱い

Last updated at Posted at 2024-08-31

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とベクターカラムを扱う方法について説明しました。特に、ベクターカラムのテストをサポートするためのモデル作成やデータベース設定の工夫について紹介しました。

参考記事

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0