如何测试使用寿命函数的 FastAPI 端点？

Question

有人可以告诉我如何测试使用FastAPI 的新生命周期功能的端点吗？

我正在尝试为使用生命周期函数中的资源的端点设置测试，但测试失败，因为我在生命周期函数中设置的字典未作为 FastAPI 应用程序的一部分传递给 TestClient。

我的 API 如下所示。

from fastapi import FastAPI
from contextlib import asynccontextmanager

ml_model = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    predictor = Predictor(model_version)
    ml_model["predict"] = predictor.predict_from_features
    yield
    # Clean up the ML models and release the resources
    ml_model.clear()


app = FastAPI(lifespan=lifespan)

@app.get("/prediction/")
async def get_prediction(model_input: str):
    prediction = ml_model["predict"](model_input)
    return prediction

端点的测试代码/prediction如下所示：

from fastapi.testclient import TestClient

from app.main import app

client = TestClient(app)

def test_read_prediction():
    model_input= "test"
    response = client.get(f"/prediction/?model_input={model_input}")
    assert response.status_code == 200

测试失败，并显示错误消息 KeyError: 'predict'，这表明该ml_models字典未与应用程序对象一起传递。我也尝试过使用 app.state.ml_models = {}，但这也不起作用。我将不胜感激任何帮助！

Answer 1

使用TestClient作为上下文管理器。这会触发启动/关闭事件以及寿命。

from fastapi.testclient import TestClient

from app.main import app

def test_read_prediction():
    with TestClient(app) as client:
        model_input= "test"
        response = client.get(f"/prediction/?model_input={model_input}")
        assert response.status_code == 200

这是一些有关测试启动/关闭事件的文档。它也适用于寿命。