mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
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.
This commit is contained in:
parent
999d44279d
commit
806db5663d
1 changed files with 89 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue