Search code examples
pythonpytestfastapi

How to use LifespanManager to test a reverse proxy in FastAPI (async testing)


According to FastAPI documentation I may need to use a LifespanManager. Can someone show me an example of how to use the LifespanManager in an async test? Like, with this lifespan:

    @asynccontextmanager
    async def lifespan(_app: FastAPI):
        async with httpx.AsyncClient(base_url=env.proxy_url, transport=httpx.MockTransport(dummy_response)) as client:
            yield {'client': client}
            await client.aclose()

I'm trying to test an endpoint called proxy, which works fine but I need tests for regression:

import pytest
import pytest_asyncio
from fastapi import FastAPI
from contextlib import asynccontextmanager
from asgi_lifespan import LifespanManager
import httpx
from httpx import Response
import importlib
import uvicorn

import proxy
import env

def dummy_response(_request):
    res = Response(200, content="Mock response")
    res.headers['Content-Type'] = 'text/plain; charset=utf-8'
    return res

@pytest_asyncio.fixture
async def mock_proxy():
    importlib.reload(env)
    importlib.reload(proxy)

    @asynccontextmanager
    async def lifespan(_app: FastAPI):
        async with httpx.AsyncClient(base_url=env.proxy_url, transport=httpx.MockTransport(dummy_response)) as client:
            yield {'client': client}
            await client.aclose()

    app = FastAPI(lifespan=lifespan)
    app.add_route("/proxy/path", proxy.proxy)

    async with LifespanManager(app) as manager:
        yield app

@pytest_asyncio.fixture
async def _client(mock_proxy):
    async with mock_proxy as app:
        async with httpx.AsyncClient(app=app, base_url=env.proxy_url) as client:
            yield client

@pytest.mark.anyio
async def test_proxy_get_request(_client):
    async with _client as client:
        response = await client.get(f"{env.proxy_url}/proxy/path", params={"query": "param"})
        assert response.status_code == 200

This attempt tells me

TypeError: 'FastAPI' object does not support the asynchronous context manager protocol

edit:

this code seems pretty close, but the lifespan change to state is not occurring:

import pytest
import pytest_asyncio
from fastapi import FastAPI
from contextlib import asynccontextmanager
from asgi_lifespan import LifespanManager
import httpx
from httpx import Response
import importlib
import uvicorn

import proxy
import env

def dummy_response(_request):
    res = Response(200, content="Mock response")
    res.headers['Content-Type'] = 'text/plain; charset=utf-8'
    return res

@pytest_asyncio.fixture
async def _client():
    with pytest.MonkeyPatch.context() as monkeypatch:
        monkeypatch.setenv("PROXY_URL", "http://proxy")

        importlib.reload(env)
        importlib.reload(proxy)

    @asynccontextmanager
    async def lifespan(_app: FastAPI):
        async with httpx.AsyncClient(
             base_url=env.proxy_url, 
             transport=httpx.MockTransport(dummy_response)) as client:
            yield {'client': client} # startup
            await client.aclose() # shutdown

    app = FastAPI(lifespan=lifespan)
    app.add_route("/proxy/path", proxy.proxy)

    transport = httpx.ASGITransport(app=app)
    async with httpx.AsyncClient(transport=transport, base_url=env.proxy_url) \
          as client, LifespanManager(app):
        yield client

@pytest.mark.asyncio
async def test_proxy_get_request(_client):
        response = await _client.get(f"/proxy/path", params={"query": "param"})
        assert response.status_code == 200

==================================================== short test summary info ==================================================== FAILED tests/regression/test_proxy.py::test_proxy_get_request - AttributeError: 'State' object has no attribute 'client'

... in fact it seems like the LifespanManager is not doing a lot of work. If I change the last part of the fixture to:

    async with LifespanManager(app) as manager:
        print(manager._state, app.state._state)
        app.state = State(state=manager._state)
        transport = httpx.ASGITransport(app=app)
        print(manager._state, app.state.client)
        async with httpx.AsyncClient(transport=transport, base_url=env.proxy_url) \
              as client:
            yield client

I get:

----------------------------------------------------- Captured stdout setup ----------------------------------------------------- {'client': <httpx.AsyncClient object at 0x1034da240>} {} {'client': <httpx.AsyncClient object at 0x1034da240>} <httpx.AsyncClient object at 0x1034da240> ==================================================== short test summary info ==================================================== FAILED tests/regression/test_proxy.py::test_proxy_get_request - AttributeError: 'State' object has no attribute 'client'

So, app is not getting state at startup (but the manager is, it's just not applied to app for some reason). Likewise, even manually setting state myself (so why even use LifespanManager at that point), the state is not available in the proxy function's request like it is supposed to be.

The reason I am doing this is the first line in the proxy is:

async def proxy(request: Request):
    client = request.state.client

And this is what is failing.

edit 2:

thanks to Yurii's comments, I resolved this initial issue, but not what led me down this path in the first place. I can get past this issue with:

    async with LifespanManager(app) as manager:
        async with httpx.AsyncClient(transport=httpx.ASGITransport(app=manager.app), base_url=env.proxy_url) as client:
                yield client

However, this all started when my initial approach with FastAPI's TestClient was failing because of a weird cancellation, triggering unhandled issues in a taskgroup, triggering the stream to be exhausted and then attempted to read (which, if there was a streaming issue, the proxy wouldn't work, and it does). It turns out FastAPI doesn't allow you to use TestClient for async tests, and recommends this approach (see the link at the start of this for more). I am now getting much the same issue here:

FAILED tests/regression/test_proxy.py::test_proxy_get_request - ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)

which is caused by the same cancellation issue:

self = <asyncio.locks.Event object at 0x105cdade0 [unset]>

async def wait(self):
    """Block until the internal flag is true.

    If the internal flag is true on entry, return True
    immediately.  Otherwise, block until another coroutine calls
    set() to set the flag to true, then return True.
    """
    if self._value:
        return True

    fut = self._get_loop().create_future()
    self._waiters.append(fut)
    try:
      await fut

E asyncio.exceptions.CancelledError: Cancelled by cancel scope 105cdb8f0


Solution

    1. You should yield manager.app from mock_proxy, not just app

    2. In _client fixture first context manager async with mock_proxy as app: and just assign mock_proxy to app instead (app = mock_proxy)

    3. You don't need context manager in test_proxy_get_request as well (remove async with _client as client: and use _client)