How to test custom middleware in FastAPI?

Sooner or later, you need to add custom middleware to your FastAPI app. For example, let's say that we want to add a request ID to every request for easier traceability. Once we have it, we'd like to test it. To make sure it's working as expected. The question is how to do that efficiently.

First instinct

So let's take a look at a real example. Let's stick with "add request ID to every request". Middleware for that can look like this:

# request_id_middleware.py
import uuid

from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response


class RequestIDMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, header_name: str = "X-Request-ID") -> None:
        super().__init__(app)
        self.header_name = header_name

    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        incoming = request.headers.get(self.header_name)
        request_id = incoming if incoming else str(uuid.uuid4())
        request.state.request_id = request_id

        response = await call_next(request)
        response.headers[self.header_name] = request_id
        return response

This middleware tries to read the request ID from the incoming request's header. If it's not there, it generates one. Then it sets it as part of the request's state. Finally, it adds the request ID value to the response header. In other words, nothing special.

The first instinct to test it would be the following:

Initialize the RequestIDMiddleware instance
Set up test input data
Call the dispatch method
Assert the header in the returned response has the expected value

The tests would then look like something like this:

# tests/test_request_id_middleware.py
import uuid
from unittest.mock import AsyncMock, MagicMock

import pytest

from request_id_middleware import RequestIDMiddleware


def make_request(headers: dict[str, str] | None = None) -> MagicMock:
    request = MagicMock()
    request.headers = headers or {}
    request.state = MagicMock()  # request.state.request_id = ... just records the attr
    return request


@pytest.mark.asyncio
async def test_generates_id_when_absent():
    mw = RequestIDMiddleware(app=MagicMock())
    request = make_request()

    response = MagicMock()
    response.headers = {}
    call_next = AsyncMock(return_value=response)

    result = await mw.dispatch(request, call_next)

    # call_next was awaited with the request
    call_next.assert_awaited_once_with(request)

    # a valid UUID was generated and stamped on both state and response
    generated = request.state.request_id
    uuid.UUID(generated)  # raises if not a valid UUID
    assert response.headers["X-Request-ID"] == generated
    assert result is response


@pytest.mark.asyncio
async def test_reuses_incoming_id():
    mw = RequestIDMiddleware(app=MagicMock())
    incoming = "abc-123-fixed"
    request = make_request({"X-Request-ID": incoming})

    response = MagicMock()
    response.headers = {}
    call_next = AsyncMock(return_value=response)

    await mw.dispatch(request, call_next)

    assert request.state.request_id == incoming
    assert response.headers["X-Request-ID"] == incoming


@pytest.mark.asyncio
async def test_custom_header_name():
    mw = RequestIDMiddleware(app=MagicMock(), header_name="X-Trace-Id")
    request = make_request({"X-Trace-Id": "trace-9"})

    response = MagicMock()
    response.headers = {}
    call_next = AsyncMock(return_value=response)

    await mw.dispatch(request, call_next)

    assert response.headers["X-Trace-Id"] == "trace-9"

The problem there is that we're using mocks all over the place. We mock the request and we mock the call_next. The big problem with that is that Request or RequestResponseEndpoint might change over time. This won't be reflected in our tests. The middleware might not be compatible with the new signature anymore, but the tests will still pass due to the passed mocks. This feels like testing that's too close to implementation details. Also, the tests look quite ugly when so much mocking is in place.

Become a better engineer, one article at a time.

Practices, mindsets, and habits that actually move the needle. Delivered weekly to your inbox.

Taking a step back

Let's take a step back. What we're really interested in is that the response header contains the request ID header and that it's propagated from the request when present. This doesn't tell us anything about the middleware internals. So let's focus on that.

Instead of testing middleware directly, we can test it indirectly. We can spin up the bare-bones FastAPI app inside for tests and add middleware to it. Then we can simply use TestClient for our tests.

import uuid
from typing import Callable

import pytest
from fastapi import FastAPI, Request
from fastapi.testclient import TestClient

from request_id_middleware import RequestIDMiddleware


@pytest.fixture
def create_client() -> Callable[..., TestClient]:
    def _create(header_name: str | None = None) -> TestClient:
        app = FastAPI()
        if header_name is None:
            app.add_middleware(RequestIDMiddleware)
        else:
            app.add_middleware(RequestIDMiddleware, header_name=header_name)

        @app.get("/")
        async def root(request: Request):
            return {"seen_id": request.state.request_id}

        return TestClient(app)

    return _create


def test_generates_id_when_absent(create_client):
    client = create_client()
    r = client.get("/")

    assert r.status_code == 200
    header_id = r.headers["X-Request-ID"]
    uuid.UUID(header_id)  # raises if not a valid UUID
    assert r.json()["seen_id"] == header_id


def test_reuses_incoming_id(create_client):
    client = create_client()
    r = client.get("/", headers={"X-Request-ID": "incoming-42"})

    assert r.headers["X-Request-ID"] == "incoming-42"
    assert r.json()["seen_id"] == "incoming-42"


def test_custom_header_name(create_client):
    client = create_client(header_name="X-Trace-Id")
    r = client.get("/", headers={"X-Trace-Id": "trace-9"})

    assert r.headers["X-Trace-Id"] == "trace-9"
    assert r.json()["seen_id"] == "trace-9"


def test_header_is_case_insensitive(create_client):
    client = create_client()
    r = client.get("/", headers={"x-request-id": "lower-case-id"})
    assert r.json()["seen_id"] == "lower-case-id"

As you can see, there's no need for mocking anymore. Test setup is much shorter. We're also sure that if anything changes around the expected middleware signature, the tests will fail as the request actually goes through the app.

Conclusion

As you can see, first instinct is not necessarily the best approach. Many times it helps to take a step back and try to see a bit bigger picture. Some might argue that now the tests are integration tests and not unit tests. It doesn't matter how you call the tests. What matters is that they are valuable.

Happy engineering!

First instinct

Become a better engineer, one article at a time.

Taking a step back

Conclusion

Stay Updated