File size: 4,265 Bytes
4aecccf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """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
|