test_testclient.py

import itertools

import anyio
import pytest

from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.responses import JSONResponse
from starlette.websockets import WebSocket, WebSocketDisconnect

mock_service = Starlette()


@mock_service.route("/")
def mock_service_endpoint(request):
    return JSONResponse({"mock": "example"})


_identity_runvar: anyio.lowlevel.RunVar[int] = anyio.lowlevel.RunVar("_identity_runvar")


def get_identity(counter):
    try:
        return _identity_runvar.get()
    except LookupError:
        token = next(counter)
        _identity_runvar.set(token)
        return token


def create_app(test_client_factory, counter=itertools.count()):
    app = Starlette()

    @app.on_event("startup")
    async def get_startup_thread():
        app.startup_task = anyio.get_current_task().id
        app.startup_loop = get_identity(counter)

    @app.on_event("shutdown")
    async def get_shutdown_thread():
        app.shutdown_task = anyio.get_current_task().id
        app.shutdown_loop = get_identity(counter)

    @app.route("/")
    def homepage(request):
        client = test_client_factory(mock_service)
        response = client.get("/")
        return JSONResponse(response.json())

    @app.route("/thread")
    async def thread(request):
        return JSONResponse(get_identity(counter))

    return app


startup_error_app = Starlette()


@startup_error_app.on_event("startup")
def startup():
    raise RuntimeError()


def test_use_testclient_in_endpoint(test_client_factory):
    """
    We should be able to use the test client within applications.

    This is useful if we need to mock out other services,
    during tests or in development.
    """
    client = test_client_factory(create_app(test_client_factory))
    response = client.get("/")
    assert response.json() == {"mock": "example"}


def test_use_testclient_as_contextmanager(test_client_factory, anyio_backend_name):
    """
    This test asserts a number of properties that are important for an
    app level task_group
    """
    app = create_app(test_client_factory, counter=itertools.count())
    client = test_client_factory(app)

    with client:
        # within a TestClient context every async request runs in the same thread
        assert client.get("/thread").json() == 0
        assert client.get("/thread").json() == 0

    # that thread is also the same as the lifespan thread
    assert app.startup_loop == 0
    assert app.shutdown_loop == 0

    # lifespan events run in the same task, this is important because a task
    # group must be entered and exited in the same task.
    assert app.startup_task == app.shutdown_task

    # outside the TestClient context, new requests continue to spawn in new
    # eventloops in new threads
    assert client.get("/thread").json() == 1
    assert client.get("/thread").json() == 2

    first_task = app.startup_task

    with client:
        # the TestClient context can be re-used, starting a new lifespan task
        # in a new thread
        assert client.get("/thread").json() == 3
        assert client.get("/thread").json() == 3

    assert app.startup_loop == 3
    assert app.shutdown_loop == 3

    # lifespan events still run in the same task, with the context but...
    assert app.startup_task == app.shutdown_task

    if anyio_backend_name != "asyncio":
        # https://github.com/agronholm/anyio/issues/324
        # ... the second TestClient context creates a new lifespan task.
        assert first_task != app.startup_task


def test_error_on_startup(test_client_factory):
    with pytest.raises(RuntimeError):
        with test_client_factory(startup_error_app):
            pass  # pragma: no cover


def test_exception_in_middleware(test_client_factory):
    class MiddlewareException(Exception):
        pass

    class BrokenMiddleware:
        def __init__(self, app):
            self.app = app

        async def __call__(self, scope, receive, send):
            raise MiddlewareException()

    broken_middleware = Starlette(middleware=[Middleware(BrokenMiddleware)])

    with pytest.raises(MiddlewareException):
        with test_client_factory(broken_middleware):
            pass  # pragma: no cover


def test_testclient_asgi2(test_client_factory):
    def app(scope):
        async def inner(receive, send):
            await send(
                {
                    "type": "http.response.start",
                    "status": 200,
                    "headers": [[b"content-type", b"text/plain"]],
                }
            )
            await send({"type": "http.response.body", "body": b"Hello, world!"})

        return inner

    client = test_client_factory(app)
    response = client.get("/")
    assert response.text == "Hello, world!"


def test_testclient_asgi3(test_client_factory):
    async def app(scope, receive, send):
        await send(
            {
                "type": "http.response.start",
                "status": 200,
                "headers": [[b"content-type", b"text/plain"]],
            }
        )
        await send({"type": "http.response.body", "body": b"Hello, world!"})

    client = test_client_factory(app)
    response = client.get("/")
    assert response.text == "Hello, world!"


def test_websocket_blocking_receive(test_client_factory):
    def app(scope):
        async def respond(websocket):
            await websocket.send_json({"message": "test"})

        async def asgi(receive, send):
            websocket = WebSocket(scope, receive=receive, send=send)
            await websocket.accept()
            async with anyio.create_task_group() as task_group:
                task_group.start_soon(respond, websocket)
                try:
                    # this will block as the client does not send us data
                    # it should not prevent `respond` from executing though
                    await websocket.receive_json()
                except WebSocketDisconnect:
                    pass

        return asgi

    client = test_client_factory(app)
    with client.websocket_connect("/") as websocket:
        data = websocket.receive_json()
        assert data == {"message": "test"}