clean up mock transport: remove streaming, add defensive parsing

litellm/llms/custom_httpx/mock_transport.py

@@ -8,7 +8,8 @@ so the full proxy -> router -> OpenAI SDK -> httpx path is exercised.
 
 import json
 import time
-from typing import Iterator, List
+import uuid
+from typing import Tuple
 
 import httpx
 
@@ -17,10 +18,14 @@ import httpx
 # Pre-built response templates
 # ---------------------------------------------------------------------------
 
+def _mock_id() -> str:
+    return f"chatcmpl-mock-{uuid.uuid4().hex[:8]}"
+
+
 def _chat_completion_json(model: str) -> dict:
     """Return a minimal valid ChatCompletion object."""
     return {
-        "id": "chatcmpl-mock",
+        "id": _mock_id(),
         "object": "chat.completion",
         "created": int(time.time()),
         "model": model,
@@ -42,74 +47,10 @@ def _chat_completion_json(model: str) -> dict:
     }
 
 
-def _streaming_sse_payloads(model: str) -> List[bytes]:
-    """Pre-build the SSE byte payloads for a streaming response."""
-    chunk = {
-        "id": "chatcmpl-mock",
-        "object": "chat.completion.chunk",
-        "created": int(time.time()),
-        "model": model,
-        "choices": [
-            {
-                "index": 0,
-                "delta": {"role": "assistant", "content": "Mock response"},
-                "finish_reason": None,
-            }
-        ],
-    }
-    done_chunk = {
-        "id": "chatcmpl-mock",
-        "object": "chat.completion.chunk",
-        "created": int(time.time()),
-        "model": model,
-        "choices": [
-            {
-                "index": 0,
-                "delta": {},
-                "finish_reason": "stop",
-            }
-        ],
-    }
-    return [
-        b"data: " + json.dumps(chunk).encode() + b"\n\n",
-        b"data: " + json.dumps(done_chunk).encode() + b"\n\n",
-        b"data: [DONE]\n\n",
-    ]
-
-
-# ---------------------------------------------------------------------------
-# Byte-stream wrappers
-# ---------------------------------------------------------------------------
-
-class MockSSEAsyncStream(httpx.AsyncByteStream):
-    """Async byte stream that yields pre-built SSE payloads."""
-
-    def __init__(self, payloads: List[bytes]) -> None:
-        self._payloads = payloads
-
-    async def __aiter__(self):  # type: ignore[override]
-        for payload in self._payloads:
-            yield payload
-
-
-class MockSSESyncStream(httpx.SyncByteStream):
-    """Sync byte stream that yields pre-built SSE payloads."""
-
-    def __init__(self, payloads: List[bytes]) -> None:
-        self._payloads = payloads
-
-    def __iter__(self) -> Iterator[bytes]:
-        return iter(self._payloads)
-
-
 # ---------------------------------------------------------------------------
 # Transport
 # ---------------------------------------------------------------------------
 
-_STREAM_HEADERS = {
-    "content-type": "text/event-stream",
-}
-
 _JSON_HEADERS = {
     "content-type": "application/json",
 }
@@ -123,22 +64,17 @@ class MockOpenAITransport(httpx.AsyncBaseTransport, httpx.BaseTransport):
     """
 
     @staticmethod
-    def _parse_request(request: httpx.Request) -> tuple:
-        """Extract (model, stream) from the request body."""
-        body = json.loads(request.content)
+    def _parse_request(request: httpx.Request) -> Tuple[str, bool]:
+        """Extract model from the request body."""
+        try:
+            body = json.loads(request.content)
+        except (json.JSONDecodeError, ValueError):
+            return ("mock-model", False)
         model = body.get("model", "mock-model")
-        stream = body.get("stream", False)
-        return model, stream
+        return (model, False)
 
     async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
-        model, stream = self._parse_request(request)
-        if stream:
-            payloads = _streaming_sse_payloads(model)
-            return httpx.Response(
-                status_code=200,
-                headers=_STREAM_HEADERS,
-                stream=MockSSEAsyncStream(payloads),
-            )
+        model, _ = self._parse_request(request)
         body = json.dumps(_chat_completion_json(model)).encode()
         return httpx.Response(
             status_code=200,
@@ -147,14 +83,7 @@ class MockOpenAITransport(httpx.AsyncBaseTransport, httpx.BaseTransport):
         )
 
     def handle_request(self, request: httpx.Request) -> httpx.Response:
-        model, stream = self._parse_request(request)
-        if stream:
-            payloads = _streaming_sse_payloads(model)
-            return httpx.Response(
-                status_code=200,
-                headers=_STREAM_HEADERS,
-                stream=MockSSESyncStream(payloads),
-            )
+        model, _ = self._parse_request(request)
         body = json.dumps(_chat_completion_json(model)).encode()
         return httpx.Response(
             status_code=200,

litellm/llms/openai/common_utils.py

@@ -22,6 +22,7 @@ from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     get_ssl_configuration,
 )
+from litellm.llms.custom_httpx.mock_transport import MockOpenAITransport
 
 
 def _get_client_init_params(cls: type) -> Tuple[str, ...]:
@@ -206,8 +207,6 @@ class BaseOpenAILLM:
             return litellm.aclient_session
 
         if getattr(litellm, "network_mock", False):
-            from litellm.llms.custom_httpx.mock_transport import MockOpenAITransport
-
             return httpx.AsyncClient(transport=MockOpenAITransport())
 
         # Get unified SSL configuration
@@ -231,8 +230,6 @@ class BaseOpenAILLM:
             return litellm.client_session
 
         if getattr(litellm, "network_mock", False):
-            from litellm.llms.custom_httpx.mock_transport import MockOpenAITransport
-
             return httpx.Client(transport=MockOpenAITransport())
 
         # Get unified SSL configuration

tests/test_litellm/llms/custom_httpx/test_mock_transport.py

@@ -1,6 +1,6 @@
 """
 Tests for MockOpenAITransport — verifies that the mock transport produces
-responses parseable by the OpenAI SDK for both streaming and non-streaming paths.
+responses parseable by the OpenAI SDK.
 """
 
 import json
@@ -60,71 +60,28 @@ class TestNonStreaming:
         body = json.loads(response.content)
         assert body["model"] == "my-custom-model"
 
-
-# ---------------------------------------------------------------------------
-# Streaming
-# ---------------------------------------------------------------------------
-
-
-class TestStreaming:
-    def test_sync_streaming_returns_sse_events(self):
+    def test_unique_ids_per_response(self):
         transport = MockOpenAITransport()
         request = httpx.Request(
             method="POST",
             url="https://api.openai.com/v1/chat/completions",
-            content=json.dumps({"model": "gpt-4o", "stream": True, "messages": []}),
+            content=json.dumps({"model": "gpt-4o", "messages": []}),
+        )
+        r1 = json.loads(transport.handle_request(request).content)
+        r2 = json.loads(transport.handle_request(request).content)
+        assert r1["id"] != r2["id"]
+
+    def test_empty_body_does_not_crash(self):
+        transport = MockOpenAITransport()
+        request = httpx.Request(
+            method="GET",
+            url="https://api.openai.com/v1/models",
+            content=b"",
         )
         response = transport.handle_request(request)
         assert response.status_code == 200
-        assert "text/event-stream" in response.headers["content-type"]
-
-        chunks = list(response.stream)
-        # Should have: content chunk, finish chunk, [DONE]
-        assert len(chunks) == 3
-        assert chunks[-1] == b"data: [DONE]\n\n"
-
-        # Parse the first chunk
-        first_line = chunks[0].decode()
-        assert first_line.startswith("data: ")
-        data = json.loads(first_line[len("data: "):].strip())
-        assert data["object"] == "chat.completion.chunk"
-        assert data["model"] == "gpt-4o"
-        assert data["choices"][0]["delta"]["content"] == "Mock response"
-
-    @pytest.mark.asyncio
-    async def test_async_streaming_returns_sse_events(self):
-        transport = MockOpenAITransport()
-        request = httpx.Request(
-            method="POST",
-            url="https://api.openai.com/v1/chat/completions",
-            content=json.dumps({"model": "gpt-4o", "stream": True, "messages": []}),
-        )
-        response = await transport.handle_async_request(request)
-        assert response.status_code == 200
-
-        chunks = []
-        async for chunk in response.stream:
-            chunks.append(chunk)
-
-        assert len(chunks) == 3
-        assert chunks[-1] == b"data: [DONE]\n\n"
-
-        # Parse finish chunk
-        finish_line = chunks[1].decode()
-        data = json.loads(finish_line[len("data: "):].strip())
-        assert data["choices"][0]["finish_reason"] == "stop"
-
-    def test_streaming_model_echoed(self):
-        transport = MockOpenAITransport()
-        request = httpx.Request(
-            method="POST",
-            url="https://api.openai.com/v1/chat/completions",
-            content=json.dumps({"model": "custom-stream", "stream": True, "messages": []}),
-        )
-        response = transport.handle_request(request)
-        first_chunk = next(iter(response.stream))
-        data = json.loads(first_chunk.decode()[len("data: "):].strip())
-        assert data["model"] == "custom-stream"
+        body = json.loads(response.content)
+        assert body["model"] == "mock-model"
 
 
 # ---------------------------------------------------------------------------