FastapiのテストでAsyncClientを使った際にlifespan eventsが走らない問題を解決する

Python

FastAPI

Kazuki Moriyama (森山和樹)

lifespanとは

fastapiの起動前、起動後に何らかの処理を仕組みたいときに利用する機構。
例えばロードに時間がかかる機械学習モデルの用意などに使える。

from contextlib import asynccontextmanager

from fastapi import FastAPI

def fake_answer_to_everything_ml_model(x: float):
    return x * 42

ml_models = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    # appが起動する前に機械学習モデルをロードする
    ml_models["answer_to_everything"] = fake_answer_to_everything_ml_model

    # yieldで実際のappが走り出す
    yield

    # appが終了しするタイミングで処理がこの関数に戻るので、
    # 機械学習モデル関連のリソースの開放
    ml_models.clear()

app = FastAPI(lifespan=lifespan)

@app.get("/predict")
async def predict(x: float):
    # 実際にモデルを使用する処理
    result = ml_models["answer_to_everything"](x)
    return {"result": result}

他にもロガーのセットアップなど、起動前後に行いたい処理を良く行う。

startup/shutdown event

lifespanは昔はstartup/shutdown eventという概念で実装されていた。

from fastapi import FastAPI

app = FastAPI()

items = {}

@app.on_event("startup")
async def startup_event():
    items["foo"] = {"name": "Fighters"}
    items["bar"] = {"name": "Tenders"}

@app.on_event("shutdown")
def shutdown_event():
    with open("log.txt", mode="a") as log:
        log.write("Application shutdown")

@app.get("/items/{item_id}")
async def read_items(item_id: str):
    return items[item_id]

ただしこれは現在は非推奨で、代わりにlifespanを使用すると良い。

この記事ではlifespan、startup/shutdown eventのどちらでも突き当たる問題について解説する。

AsyncClient

テストを書くときにたまにテストの関数自体をasyncにしたいときがある。
fastapiでは通常apiのテストを書くときにはTestClientを使用するが、asyncな関数の中ではTestClientは使用できない。
そういうときにはhttpxが用意しているAsyncClientを使用してテストを書く。

import pytest
from httpx import AsyncClient

from .main import app

@pytest.mark.anyio
async def test_root():
    async with AsyncClient(app=app, base_url="http://test") as ac:
        response = await ac.get("/")
    assert response.status_code == 200
    assert response.json() == {"message": "Tomato"}

こうすればテストの関数をasyncにしつつ、fastapiにリクエストを送信するテストが書ける。

同期的なテスト関数でのfastapiの起動前後処理の扱い

同期的なテスト関数では普通にTestClientが使用できて、TestClientを使用しているときには起動前後処理は実際のapp起動時と同様に処理される。

from fastapi import FastAPI
from fastapi.testclient import TestClient

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 起動前後処理
    ...

app = FastAPI(lifespan=lifespan)

items = {}

@app.get("/items/{item_id}")
async def read_items(item_id: str):
    return items[item_id]

def test_read_items():
    # withでTestClientを使用すれば起動前後の処理が走る
    with TestClient(app) as client:
        response = client.get("/items/foo")
        assert response.status_code == 200
        assert response.json() == {"name": "Fighters"}

AsyncClientを使用した際の問題

AsyncClientを使用してasyncな関数でテストを書いた場合には、そのままでは起動前後処理は走らない。
これはhttpxのデザインによるものだそうで、そもそも｀httpx`は単なるhttp requestライブラリだからASGIサーバーの起動前後処理などは考慮に入れたくないらしい。

とはいっても、fastapiのテストで使用する際には困るのでasgi-lifespanを使えば、無理やりlifespanを埋め込める。

import pytest
from asgi_lifespan import LifespanManager
from httpx import AsyncClient

from .main import app

@pytest.mark.anyio
async def test_root():
    # ここでlifespanを読み込ませる
    async with AsyncClient(app=app, base_url="http://test") as ac, LifespanManager(app):
        response = await ac.get("/")
    assert response.status_code == 200
    assert response.json() == {"message": "Tomato"}

あとはこの処理を何度も書くのが面倒なので、pytestを使用していれば以下のようなfixtureをどこかで定義して使い回せば良い。

# どこかのconftest.py
@pytest.fixture()
async def api_client():
    async with AsyncClient(app=app, base_url="http://test") as client, LifespanManager(app):
        yield client

import pytest
from asgi_lifespan import LifespanManager
from httpx import AsyncClient

from .main import app

@pytest.mark.anyio
# fixtureとして利用する
async def test_root(api_client): 
    response = await api_client.get("/")
    assert response.status_code == 200
    assert response.json() == {"message": "Tomato"}