From 806db5663d3c8e5c6d6caf2a98f339779e101815 Mon Sep 17 00:00:00 2001
From: Nathan Price <njprice2@illinois.edu>
Date: Sat, 13 Jun 2026 07:35:21 -0500
Subject: [PATCH] test(router): add coverage for ttft_timeout edge cases

Add tests for the empty-stream path (StopAsyncIteration -> APIError)
and the _acompletion intercept path that forces stream=True internally
when ttft_timeout is set.
---
 tests/test_litellm/test_router.py | 89 +++++++++++++++++++++++++++++++
 1 file changed, 89 insertions(+)

diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index a10e6fb7d50..ea7ff13b96e 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -4890,3 +4890,92 @@ async def test_router_ttft_timeout_not_reset_by_preamble_chunks():
         )
 
     assert "ttft_timeout" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_router_ttft_timeout_empty_stream_raises_api_error():
+    """Provider returns an empty stream (StopAsyncIteration immediately); stream_chunk_builder
+    returns None for an empty chunk list, which should raise litellm.APIError."""
+    import litellm
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "test-model",
+                "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake-key"},
+            }
+        ],
+        ttft_timeout=5.0,
+    )
+
+    async def empty_stream():
+        return
+        yield  # make it an async generator
+
+    fake_stream = MagicMock()
+    fake_stream.model = "gpt-4o"
+    fake_stream.custom_llm_provider = "openai"
+    fake_stream.__aiter__ = lambda self: empty_stream()
+
+    with pytest.raises(litellm.APIError):
+        await router._collect_stream_with_ttft_timeout(
+            response=fake_stream,
+            messages=[{"role": "user", "content": "hi"}],
+            ttft_timeout=5.0,
+        )
+
+
+@pytest.mark.asyncio
+async def test_router_ttft_timeout_acompletion_intercept():
+    """When ttft_timeout is set and stream=False, _acompletion forces stream=True internally
+    and routes the response through _collect_stream_with_ttft_timeout."""
+    from unittest.mock import AsyncMock, patch
+
+    import litellm
+    from litellm import ModelResponse
+    from litellm.utils import CustomStreamWrapper
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "test-model",
+                "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake-key"},
+            }
+        ],
+        ttft_timeout=5.0,
+    )
+
+    reconstructed = MagicMock(spec=ModelResponse)
+
+    fake_stream = MagicMock(spec=CustomStreamWrapper)
+    fake_stream.model = "gpt-4o"
+    fake_stream.custom_llm_provider = "openai"
+
+    with (
+        patch.object(
+            router,
+            "_collect_stream_with_ttft_timeout",
+            new_callable=AsyncMock,
+            return_value=reconstructed,
+        ) as mock_collect,
+        patch.object(router, "async_get_available_deployment") as mock_dep,
+        patch.object(router, "_update_kwargs_with_deployment"),
+        patch.object(router, "async_routing_strategy_pre_call_checks"),
+        patch.object(router, "_get_client", return_value=None),
+        patch.object(router, "_track_deployment_metrics"),
+        patch.object(router, "_should_raise_content_policy_error", return_value=False),
+        patch("litellm.acompletion", new_callable=AsyncMock, return_value=fake_stream),
+    ):
+        mock_dep.return_value = {
+            "model_name": "test-model",
+            "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake-key"},
+        }
+
+        result = await router._acompletion(
+            model="test-model",
+            messages=[{"role": "user", "content": "hi"}],
+            stream=False,
+        )
+
+    assert result is reconstructed
+    assert mock_collect.called