diff --git a/tests/integration/test_otel_trace_propagation.py b/tests/integration/test_otel_trace_propagation.py new file mode 100644 index 000000000..11c0b41a4 --- /dev/null +++ b/tests/integration/test_otel_trace_propagation.py @@ -0,0 +1,241 @@ +"""Integration tests for OpenTelemetry trace context propagation. + +Verifies that trace context is correctly propagated across service +boundaries and that spans share trace IDs with correct parent-child +relationships when flowing through the query endpoint and its +downstream components. +""" + +# pylint: disable=protected-access + +from collections.abc import Generator + +import pytest +from fastapi import Request +from opentelemetry import context as otel_context +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + +from app.endpoints.query import query_endpoint_handler +from authentication.interface import AuthTuple +from models.api.requests import QueryRequest + +KNOWN_TRACE_ID = "4bf92f3577b34da6a3ce929d0e0e4736" +KNOWN_PARENT_SPAN_ID = "00f067aa0ba902b7" +TRACEPARENT = f"00-{KNOWN_TRACE_ID}-{KNOWN_PARENT_SPAN_ID}-01" + + +@pytest.fixture(name="otel_collector", scope="module") +def otel_collector_fixture() -> Generator[InMemorySpanExporter, None, None]: + """Install a global TracerProvider backed by an InMemorySpanExporter. + + Module-scoped so that every module-level ``trace.get_tracer(__name__)`` + proxy resolves to the same real tracer across all tests in this file. + Shutting down the provider between tests would invalidate the cached + proxy and silently drop spans. + + Why we reset ``_TRACER_PROVIDER_SET_ONCE`` instead of using + ``mock.patch("opentelemetry.trace.get_tracer_provider")``: + ``ProxyTracer._tracer`` reads the module-level ``_TRACER_PROVIDER`` + variable directly — it never calls ``get_tracer_provider()``. + Patching the function leaves that variable ``None``, so the proxy + falls back to a noop tracer and no spans are captured. + ``set_tracer_provider()`` is the only way to populate the variable, + and the one-shot guard must be reset to allow re-installation. + """ + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + trace._TRACER_PROVIDER_SET_ONCE._done = False + trace._TRACER_PROVIDER = None + trace.set_tracer_provider(provider) + + yield exporter + + provider.shutdown() + trace._TRACER_PROVIDER_SET_ONCE._done = False + trace._TRACER_PROVIDER = None + + +@pytest.fixture(autouse=True) +def _clear_spans(otel_collector: InMemorySpanExporter) -> None: + """Clear collected spans before each test.""" + otel_collector.clear() + + +def _inject_w3c_context(traceparent: str) -> object: + """Extract a W3C traceparent header into OTel context and attach it. + + Parameters: + traceparent: W3C Trace Context header value. + + Returns: + Context token to pass to ``otel_context.detach``. + """ + ctx = TraceContextTextMapPropagator().extract({"traceparent": traceparent}) + return otel_context.attach(ctx) + + +# ============================================================================ +# Tests +# ============================================================================ + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("test_config", "mock_ogx_client", "mock_query_agent") +async def test_incoming_trace_context_is_continued( + mock_request_with_auth: Request, + test_auth: AuthTuple, + otel_collector: InMemorySpanExporter, +) -> None: + """Spans continue the trace ID received in a W3C traceparent header.""" + token = _inject_w3c_context(TRACEPARENT) + try: + await query_endpoint_handler( + request=mock_request_with_auth, + query_request=QueryRequest( # pyright: ignore[reportCallIssue] + query="What is Ansible?" + ), + auth=test_auth, + mcp_headers={}, + ) + finally: + otel_context.detach(token) # pyright: ignore[reportArgumentType] + + spans = otel_collector.get_finished_spans() + assert len(spans) > 0, "Expected at least one span" + + expected_trace_id = int(KNOWN_TRACE_ID, 16) + for span in spans: + assert span.context is not None + assert span.context.trace_id == expected_trace_id, ( + f"Span {span.name!r} has trace_id " + f"{span.context.trace_id:#034x}, expected {expected_trace_id:#034x}" + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("test_config", "mock_ogx_client", "mock_query_agent") +async def test_root_span_is_child_of_incoming_parent( + mock_request_with_auth: Request, + test_auth: AuthTuple, + otel_collector: InMemorySpanExporter, +) -> None: + """The endpoint root span's parent points to the incoming span ID.""" + token = _inject_w3c_context(TRACEPARENT) + try: + await query_endpoint_handler( + request=mock_request_with_auth, + query_request=QueryRequest( # pyright: ignore[reportCallIssue] + query="What is Ansible?" + ), + auth=test_auth, + mcp_headers={}, + ) + finally: + otel_context.detach(token) # pyright: ignore[reportArgumentType] + + spans = otel_collector.get_finished_spans() + root_spans = [s for s in spans if s.name == "query.handle_request"] + assert len(root_spans) == 1 + + root = root_spans[0] + assert ( + root.parent is not None + ), "Root span should be a child of the incoming context" + assert root.parent.span_id == int(KNOWN_PARENT_SPAN_ID, 16) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("test_config", "mock_ogx_client", "mock_query_agent") +async def test_spans_across_components_share_trace_id( + mock_request_with_auth: Request, + test_auth: AuthTuple, + otel_collector: InMemorySpanExporter, +) -> None: + """All spans emitted during a single request share the same trace ID.""" + await query_endpoint_handler( + request=mock_request_with_auth, + query_request=QueryRequest( # pyright: ignore[reportCallIssue] + query="What is Ansible?" + ), + auth=test_auth, + mcp_headers={}, + ) + + spans = otel_collector.get_finished_spans() + assert len(spans) > 1, "Expected spans from multiple components" + + trace_ids = {span.context.trace_id for span in spans if span.context is not None} + assert ( + len(trace_ids) == 1 + ), f"All spans must share one trace ID, got {len(trace_ids)}" + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("test_config", "mock_ogx_client", "mock_query_agent") +async def test_parent_child_relationships_preserved( + mock_request_with_auth: Request, + test_auth: AuthTuple, + otel_collector: InMemorySpanExporter, +) -> None: + """Child spans (quota, shield, RAG, inference) are parented to the root span.""" + await query_endpoint_handler( + request=mock_request_with_auth, + query_request=QueryRequest( # pyright: ignore[reportCallIssue] + query="What is Ansible?" + ), + auth=test_auth, + mcp_headers={}, + ) + + spans = otel_collector.get_finished_spans() + + root_spans = [s for s in spans if s.name == "query.handle_request"] + assert len(root_spans) == 1 + root = root_spans[0] + assert root.context is not None + + child_spans = [s for s in spans if s.name != "query.handle_request"] + assert len(child_spans) >= 1, "Expected at least one child span" + + for child in child_spans: + assert child.parent is not None, f"Span {child.name!r} should have a parent" + assert ( + child.parent.span_id == root.context.span_id + ), f"Span {child.name!r} should be parented to query.handle_request" + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("test_config", "mock_ogx_client", "mock_query_agent") +async def test_expected_child_spans_are_emitted( + mock_request_with_auth: Request, + test_auth: AuthTuple, + otel_collector: InMemorySpanExporter, +) -> None: + """The query flow emits the expected set of child spans.""" + await query_endpoint_handler( + request=mock_request_with_auth, + query_request=QueryRequest( # pyright: ignore[reportCallIssue] + query="What is Ansible?" + ), + auth=test_auth, + mcp_headers={}, + ) + + span_names = {s.name for s in otel_collector.get_finished_spans()} + + expected = { + "query.handle_request", + "quota.check", + "shield.moderate", + "llm.inference", + } + missing = expected - span_names + assert not missing, f"Missing expected spans: {missing}"