From 5274719d66b6ab0cfd418cd2eddd7673fcd2b14a Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 03:47:09 +0000 Subject: [PATCH] test(otel): cover websocket trace propagation with forwarding on and off Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../pass_through_endpoints/test_pass_through_endpoints.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index de951278487..6d941bca485 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4884,7 +4884,8 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): @pytest.mark.asyncio -async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch): +@pytest.mark.parametrize("forward_headers", [True, False]) +async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch, forward_headers: bool): from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import get_current_span from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator @@ -4931,14 +4932,14 @@ async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch target="wss://upstream.example.test/v1/realtime", custom_headers={}, user_api_key_dict=UserAPIKeyAuth(), - forward_headers=True, + forward_headers=forward_headers, endpoint="/realtime", accept_websocket=True, ) propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"])) assert propagated.get_span_context().trace_id == span.get_span_context().trace_id - assert captured["headers"]["authorization"] == "Bearer client" + assert captured["headers"].get("authorization") == ("Bearer client" if forward_headers else None) class ClosingUpstreamWebSocket: