Function test_websocket_route_middleware

tests/test_routing.py:991–1026

(
    test_client_factory: TestClientFactory,
)

def test_websocket_route_middleware(
    test_client_factory: TestClientFactory,
) -> None:
    async def websocket_endpoint(session: WebSocket) -> None:
        await session.accept()
        await session.send_text("Hello, world!")
        await session.close()

    class WebsocketMiddleware:
        def __init__(self, app: ASGIApp) -> None:
            self.app = app

        async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
            async def modified_send(msg: Message) -> None:
                if msg["type"] == "websocket.accept":
                    msg["headers"].append((b"X-Test", b"Set by middleware"))
                await send(msg)

            await self.app(scope, receive, modified_send)

    app = Starlette(
        routes=[
            WebSocketRoute(
                "/ws",
                endpoint=websocket_endpoint,
                middleware=[Middleware(WebsocketMiddleware)],
            )
        ]
    )

    client = test_client_factory(app)

    with client.websocket_connect("/ws") as websocket:
        text = websocket.receive_text()
        assert text == "Hello, world!"
        assert websocket.extra_headers == [(b"X-Test", b"Set by middleware")]


def test_route_repr() -> None: