wehub-resource-sync
141 lines
4.2 KiB
Python
141 lines
4.2 KiB
Python
"""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
|