scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 233

test(proxy): use controlled clocks for Responses keepalive failures

Browse source
This commit is contained in:
HootRock 2026-10-28 20:35:05 +08:00
parent 7586485b4e
commit 9f09276e73
2 changed files with 91 additions and 73 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

112
tests/unit/proxy/response_api_endpoints/test_endpoints.py
View file

@ -24,6 +24,7 @@ from litellm.proxy.proxy_server import app
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
from litellm.types.router import RouterErrors
from tests._master_key import MASTER_KEY
from tests._support.clock_loops import AutojumpClockLoop
@pytest.mark.asyncio
@ -177,65 +178,72 @@ async def test_streaming_upstream_errors_keep_the_client_protocol(
assert "error" in events[-1]
@pytest.mark.asyncio
@pytest.mark.parametrize("path", ["/v1/responses", "/v1/chat/completions"])
async def test_a_failure_after_keepalive_pings_keeps_the_client_protocol(monkeypatch: pytest.MonkeyPatch, path: str):
import litellm.proxy.proxy_server as ps
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
def test_a_failure_after_keepalive_pings_keeps_the_client_protocol(monkeypatch: pytest.MonkeyPatch, path: str) -> None:
async def request_and_check() -> None:
import litellm.proxy.proxy_server as ps
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
model: Final = "gpt-6-astra"
message: Final = "Upstream is rate limiting this key"
upstream_url: Final = "https://streaming.example/v1"
is_chat: Final = path == "/v1/chat/completions"
model: Final = "gpt-6-astra"
message: Final = "Upstream is rate limiting this key"
upstream_url: Final = "https://streaming.example/v1"
is_chat: Final = path == "/v1/chat/completions"
async def slow_rate_limit(_: httpx.Request) -> httpx.Response:
await asyncio.sleep(0.3)
return httpx.Response(429, json={"error": {"message": message, "type": "rate_limit_error", "code": None}})
async def slow_rate_limit(_: httpx.Request) -> httpx.Response:
await asyncio.sleep(0.3)
return httpx.Response(429, json={"error": {"message": message, "type": "rate_limit_error", "code": None}})
router: Final = litellm.Router(
model_list=[
{
"model_name": model,
"litellm_params": {
"model": "openai/" + model,
"api_base": upstream_url,
"api_key": "fixture-key",
"max_retries": 0,
},
}
],
num_retries=0,
)
monkeypatch.setattr(ps, "llm_router", router)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", 0.05)
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override)
with respx.mock as transport:
transport.post(upstream_url + ("/chat/completions" if is_chat else "/responses")).mock(
side_effect=slow_rate_limit
router: Final = litellm.Router(
model_list=[
{
"model_name": model,
"litellm_params": {
"model": "openai/" + model,
"api_base": upstream_url,
"api_key": "fixture-key",
"max_retries": 0,
},
}
],
num_retries=0,
)
async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client:
result: Final = await client.post(
path,
json={
"model": model,
"stream": True,
**({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"}),
},
monkeypatch.setattr(ps, "llm_router", router)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", 0.05)
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override)
with respx.mock as transport:
transport.post(upstream_url + ("/chat/completions" if is_chat else "/responses")).mock(
side_effect=slow_rate_limit
)
frames: Final = tuple(frame for frame in result.text.split("\n\n") if frame)
async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client:
result: Final = await client.post(
path,
json={
"model": model,
"stream": True,
**({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"}),
},
)
frames: Final = tuple(frame for frame in result.text.split("\n\n") if frame)
assert result.status_code == 200, result.text
assert frames[0] == ": ping", result.text
if is_chat:
assert frames[-1] == "data: [DONE]", result.text
assert message in json.loads(frames[-2].removeprefix("data: "))["error"]["message"]
else:
assert "data: [DONE]" not in result.text
assert frames[-1].startswith("event: response.failed\ndata: "), result.text
failed_error: Final = json.loads(frames[-1].split("data: ", 1)[1])["response"]["error"]
assert failed_error["code"] == "rate_limit_exceeded"
assert message in failed_error["message"]
assert result.status_code == 200, result.text
assert frames[0] == ": ping", result.text
if is_chat:
assert frames[-1] == "data: [DONE]", result.text
assert message in json.loads(frames[-2].removeprefix("data: "))["error"]["message"]
else:
assert "data: [DONE]" not in result.text
assert frames[-1].startswith("event: response.failed\ndata: "), result.text
failed_error: Final = json.loads(frames[-1].split("data: ", 1)[1])["response"]["error"]
assert failed_error["code"] == "rate_limit_exceeded"
assert message in failed_error["message"]
loop: Final = AutojumpClockLoop()
try:
loop.run_until_complete(request_and_check())
finally:
loop.run_until_complete(loop.shutdown_asyncgens())
loop.close()
_COMPACTION_ITEM: Final = {

52
tests/unit/proxy/test_common_request_processing.py
View file

@ -3,7 +3,7 @@ import contextvars
import copy
import datetime
import json
from collections.abc import AsyncIterable, Mapping
from collections.abc import AsyncIterable, Awaitable, Mapping
from types import MappingProxyType, SimpleNamespace
from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence
from urllib.parse import unquote_plus
@ -81,6 +81,7 @@ from litellm.router import Router
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.utils import ModelResponseListIterator
from litellm.router_utils.add_retry_fallback_headers import prepare_response_for_header_attachment
from tests._support.clock_loops import AutojumpClockLoop
def test_attach_guardrail_information_copies_recorded_entries_onto_model_response():
@ -10391,17 +10392,18 @@ async def test_ttft_keepalive_relays_a_late_non_streaming_body_as_an_sse_frame()
assert collected[-1] == b"data: [DONE]\n\n"
async def _slow_rate_limit_failure():
async def _slow_rate_limit_failure() -> Response:
await asyncio.sleep(0.2)
raise HTTPException(status_code=429, detail="rate limited")
async def _slow_guardrail_block_body():
async def _slow_guardrail_block_body() -> Response:
await asyncio.sleep(0.2)
return JSONResponse(status_code=400, content={"error": {"message": "blocked", "type": "guardrail_error", "code": "400"}})
return JSONResponse(
status_code=400, content={"error": {"message": "blocked", "type": "guardrail_error", "code": "400"}}
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"produce,expected_code,expected_message",
[
@ -10409,23 +10411,31 @@ async def _slow_guardrail_block_body():
(_slow_guardrail_block_body, "invalid_request_error", "blocked"),
],
)
async def test_ttft_keepalive_ends_a_late_responses_failure_with_response_failed_and_no_done(
produce, expected_code, expected_message
):
response: Final = await open_sse_before_first_byte(
produce(), ping_interval_seconds=0.05, late_framing=RESPONSES_LATE_SSE_FRAMING
)
collected: Final = await _drain(response)
def test_ttft_keepalive_ends_a_late_responses_failure_with_response_failed_and_no_done(
produce: Callable[[], Awaitable[Response]], expected_code: str, expected_message: str
) -> None:
async def collect_and_check() -> None:
response: Final = await open_sse_before_first_byte(
produce(), ping_interval_seconds=0.05, late_framing=RESPONSES_LATE_SSE_FRAMING
)
collected: Final = await _drain(response)
assert collected[0] == TTFT_PING
assert b"data: [DONE]\n\n" not in collected
event_line, data_line = collected[-1].decode().strip().split("\n")
failed: Final = json.loads(data_line.removeprefix("data: "))
assert event_line == "event: response.failed"
assert failed["type"] == "response.failed"
assert failed["response"]["status"] == "failed"
assert failed["response"]["error"]["code"] == expected_code
assert expected_message in failed["response"]["error"]["message"]
assert collected[0] == TTFT_PING
assert b"data: [DONE]\n\n" not in collected
event_line, data_line = collected[-1].decode().strip().split("\n")
failed: Final = json.loads(data_line.removeprefix("data: "))
assert event_line == "event: response.failed"
assert failed["type"] == "response.failed"
assert failed["response"]["status"] == "failed"
assert failed["response"]["error"]["code"] == expected_code
assert expected_message in failed["response"]["error"]["message"]
loop: Final = AutojumpClockLoop()
try:
loop.run_until_complete(collect_and_check())
finally:
loop.run_until_complete(loop.shutdown_asyncgens())
loop.close()
@pytest.mark.asyncio