diff --git a/tests/test_litellm/llms/azure/realtime/test_handler.py b/tests/test_litellm/llms/azure/realtime/test_handler.py index edf1b8b290f..e9d24b459d8 100644 --- a/tests/test_litellm/llms/azure/realtime/test_handler.py +++ b/tests/test_litellm/llms/azure/realtime/test_handler.py @@ -21,7 +21,9 @@ class _RecordingClientWebSocket: @pytest.mark.asyncio async def test_async_realtime_upstream_handshake_refusal_sends_error_event_then_policy_close(): - import websockets + from websockets.datastructures import Headers + from websockets.exceptions import InvalidStatus + from websockets.http11 import Response from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime from litellm.types.realtime import RealtimeErrorEvent @@ -32,9 +34,7 @@ async def test_async_realtime_upstream_handshake_refusal_sends_error_event_then_ dummy_websocket = _RecordingClientWebSocket() dummy_logging_obj = MagicMock() - refused = websockets.exceptions.InvalidStatus( - websockets.http11.Response(401, "Unauthorized", websockets.datastructures.Headers()) - ) + refused = InvalidStatus(Response(401, "Unauthorized", Headers())) with patch("websockets.connect", side_effect=refused): await handler.async_realtime( # pyright: ignore[reportUnknownMemberType] # handler's websocket param is a Protocol here but the mock connect type is incomplete diff --git a/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py b/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py index 7cd2b9e259c..f7a88b5ba63 100644 --- a/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py +++ b/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py @@ -422,7 +422,9 @@ async def test_async_realtime_ws_url_has_no_ssl(): async def test_async_realtime_upstream_handshake_refusal_sends_error_event_then_policy_close(): from typing import cast - import websockets + from websockets.datastructures import Headers + from websockets.exceptions import InvalidStatus + from websockets.http11 import Response from litellm.llms.openai.realtime.handler import OpenAIRealtime from litellm.types.realtime import RealtimeErrorEvent @@ -445,9 +447,7 @@ async def test_async_realtime_upstream_handshake_refusal_sends_error_event_then_ dummy_websocket = RecordingClientWebSocket() dummy_logging_obj = MagicMock() - refused = websockets.exceptions.InvalidStatus( - websockets.http11.Response(401, "Unauthorized", websockets.datastructures.Headers()) - ) + refused = InvalidStatus(Response(401, "Unauthorized", Headers())) with patch("websockets.connect", side_effect=refused): await handler.async_realtime( # pyright: ignore[reportUnknownMemberType] # handler's websocket param is Any diff --git a/tests/test_litellm/test_redact_string_in_error_paths.py b/tests/test_litellm/test_redact_string_in_error_paths.py index 07d1ec5f523..a5128a87b0d 100644 --- a/tests/test_litellm/test_redact_string_in_error_paths.py +++ b/tests/test_litellm/test_redact_string_in_error_paths.py @@ -2,7 +2,7 @@ Tests for _redact_string usage in error/logging paths. Covers actual execution of redaction in: -- WebSocket close reasons in realtime handlers (openai, azure, bedrock) +- WebSocket close reasons in realtime handlers (openai, bedrock) - Gemini RAG ingestion x-goog-api-key header usage - Traceback redaction pattern used in proxy streaming - Router fallback-failure traceback redaction @@ -72,25 +72,6 @@ class TestOpenAIRealtimeRedaction: api_key="test-key", ) - @pytest.mark.asyncio - async def test_invalid_status_code_redacts_reason(self): - import websockets.exceptions - - from litellm.llms.openai.realtime.handler import OpenAIRealtime - - handler = OpenAIRealtime() - exc = websockets.exceptions.InvalidStatusCode(403, None) - exc.status_code = 403 - - kwargs = self._call_kwargs() - mock_ws = kwargs["websocket"] - p1, p2, p3 = self._make_patches(handler) - with p1, p2, p3, patch("websockets.connect", side_effect=exc): - await handler.async_realtime(**kwargs) - - mock_ws.close.assert_called_once() - assert mock_ws.close.call_args[1]["code"] == 403 - @pytest.mark.asyncio async def test_generic_exception_redacts_reason(self): from litellm.llms.openai.realtime.handler import OpenAIRealtime @@ -111,41 +92,6 @@ class TestOpenAIRealtimeRedaction: assert "sk-1234567890abcdefghij" not in mock_ws.close.call_args[1]["reason"] -class TestAzureRealtimeRedaction: - """Test that Azure realtime handler redacts secrets in websocket close reasons.""" - - @pytest.mark.asyncio - async def test_invalid_status_code_redacts_reason(self): - import websockets.exceptions - - from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime - - handler = AzureOpenAIRealtime() - mock_ws = AsyncMock() - exc = websockets.exceptions.InvalidStatusCode(403, None) - exc.status_code = 403 - - with ( - patch.object( - handler, - "_construct_url", - return_value="wss://test.openai.azure.com/openai/realtime", - ), - patch("websockets.connect", side_effect=exc), - ): - await handler.async_realtime( - model="gpt-4", - websocket=mock_ws, - logging_obj=MagicMock(), - api_base="https://test.openai.azure.com/", - api_key="test-key", - api_version="2024-10-01-preview", - ) - - mock_ws.close.assert_called_once() - assert mock_ws.close.call_args[1]["code"] == 403 - - class TestBedrockRealtimeRedaction: """Test that _redact_string produces safe close reasons for Bedrock-style errors."""