How to test a FastAPI endpoint that uses lifespan function?

Question

Could someone tell me how I can test an endpoint that uses the new lifespan feature from FastAPI?

I am trying to set up tests for my endpoints that use resources from the lifespan function, but the test failed since the dict I set up in the lifespan function is not passed to the TestClient as part of the FastAPI app.

My API looks as follows.

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

And the test code for the /prediction endpoint looks as follows:

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

The test failed with an error message saying KeyError: 'predict', which shows that the ml_models dict was not passed with the app object. I also tried using app.state.ml_models = {}, but that didn't work either. I would appreciate any help!

Answer 1

Use the TestClient as a context manager. This triggers startup/shutdown events as well as lifespans.

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

Here's some documentation about testing startup/shutdown events. It also applies to lifespans.