From 7ba073aa2668820a35c69dede3ef46ed7b837cc6 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 a21f695493b..0603b5cb2ff 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 @@ -4912,7 +4912,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 @@ -4959,14 +4960,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: