"""Tests for the pure-ASGI ingress correlation owner.""" import asyncio from collections.abc import Iterator from contextlib import contextmanager from typing import cast from unittest.mock import patch import pytest from fastapi import Request from starlette.datastructures import Headers from starlette.middleware.base import BaseHTTPMiddleware from starlette.types import ASGIApp, Message, Scope from free_claude_code.api.request_ids import ( RequestCorrelationMiddleware, get_request_id, ) from tests.api.support import create_test_app def _http_scope(path: str) -> Scope: return cast( Scope, { "type": "http", "asgi": {"version": "3.0", "spec_version": "2.4"}, "http_version": "1.1", "method": "POST", "scheme": "http", "path": path, "raw_path": path.encode(), "query_string": b"", "headers": [(b"anthropic-session-id", b"session_test")], "client": None, "server": None, }, ) def test_application_uses_the_pure_asgi_correlation_owner() -> None: app = create_test_app() middleware_classes = [middleware.cls for middleware in app.user_middleware] assert sum(cls is RequestCorrelationMiddleware for cls in middleware_classes) == 1 assert all(cls is not BaseHTTPMiddleware for cls in middleware_classes) @pytest.mark.asyncio async def test_correlation_context_and_headers_span_the_complete_stream() -> None: response_started = asyncio.Event() allow_body = asyncio.Event() sent: list[Message] = [] context_entries: list[dict[str, object]] = [] context_exits: list[dict[str, object]] = [] app_request_id: str | None = None async def app(scope: Scope, _receive, send) -> None: nonlocal app_request_id app_request_id = get_request_id(Request(scope)) await send( { "type": "http.response.start", "status": 200, "headers": [], } ) response_started.set() await allow_body.wait() await send( { "type": "http.response.body", "body": b"done", "more_body": False, } ) @contextmanager def contextualize(**fields: object) -> Iterator[None]: context_entries.append(fields) try: yield finally: context_exits.append(fields) async def receive() -> Message: raise AssertionError("Test application does not receive messages") async def send(message: Message) -> None: sent.append(message) middleware = RequestCorrelationMiddleware(cast(ASGIApp, app)) with patch( "free_claude_code.api.request_ids.logger.contextualize", side_effect=contextualize, ): request = asyncio.create_task( middleware(_http_scope("/v1/responses"), receive, send) ) await response_started.wait() assert context_exits == [] assert app_request_id is not None headers = Headers(raw=sent[0]["headers"]) assert headers["request-id"] == app_request_id assert headers["x-request-id"] == app_request_id assert context_entries == [ { "http_method": "POST", "http_path": "/v1/responses", "claude_session_id": "session_test", "request_id": app_request_id, } ] allow_body.set() await request assert context_exits == context_entries @pytest.mark.asyncio async def test_correlation_middleware_passes_non_http_scopes_unchanged() -> None: observed_scope: Scope | None = None async def app(scope: Scope, _receive, _send) -> None: nonlocal observed_scope observed_scope = scope async def receive() -> Message: return {"type": "lifespan.startup"} async def send(_message: Message) -> None: return None scope = cast(Scope, {"type": "lifespan", "asgi": {"version": "3.0"}}) await RequestCorrelationMiddleware(cast(ASGIApp, app))(scope, receive, send) assert observed_scope is scope assert "state" not in scope