diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index d31ddaf0adf..4a04d33cc42 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -17,7 +17,6 @@ import respx from fastapi import HTTPException from fastapi.testclient import TestClient from httpx import Response -from starlette.types import Message, Receive, Scope, Send import litellm from litellm.proxy._types import ProxyException @@ -25,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 @@ -178,77 +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" - first_ping_sent: Final = asyncio.Event() + 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 rate_limit_after_first_ping(_: httpx.Request) -> httpx.Response: - await first_ping_sent.wait() - return httpx.Response(429, json={"error": {"message": message, "type": "rate_limit_error", "code": None}}) - - async def app_marking_first_ping(scope: Scope, receive: Receive, send: Send) -> None: - async def send_and_mark(asgi_message: Message) -> None: - await send(asgi_message) - if asgi_message["type"] == "http.response.body" and asgi_message.get("body", b"").startswith(b": ping"): - first_ping_sent.set() - - await app(scope, receive, send_and_mark) - - 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=rate_limit_after_first_ping + 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_marking_first_ping), 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: error\ndata: "), result.text - error: Final = json.loads(frames[-1].split("data: ", 1)[1])["error"] - assert error["code"] == "429" - assert message in 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: error\ndata: "), result.text + error: Final = json.loads(frames[-1].split("data: ", 1)[1])["error"] + assert error["code"] == "429" + assert message in 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 = { diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index ecc3ff860d8..944432f91b1 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -3,9 +3,9 @@ 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, Awaitable, Callable, Final, Iterator, Literal, Optional, Sequence +from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence from urllib.parse import unquote_plus from unittest.mock import AsyncMock, MagicMock, patch @@ -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(): @@ -10404,14 +10405,16 @@ 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 _rate_limit_failure_once_released(released: asyncio.Event) -> JSONResponse: - await released.wait() +async def _slow_rate_limit_failure() -> Response: + await asyncio.sleep(0.2) raise HTTPException(status_code=429, detail="rate limited") -async def _guardrail_block_body_once_released(released: asyncio.Event) -> JSONResponse: - await released.wait() - return JSONResponse(status_code=400, content={"error": {"message": "blocked", "type": "guardrail_error", "code": "400"}}) +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"}} + ) _POLICY_BLOCK_DETAIL: Final = MappingProxyType( @@ -10419,43 +10422,45 @@ _POLICY_BLOCK_DETAIL: Final = MappingProxyType( ) -async def _dict_detail_failure_once_released(released: asyncio.Event) -> JSONResponse: - await released.wait() +async def _slow_dict_detail_failure() -> Response: + await asyncio.sleep(0.2) raise HTTPException(status_code=400, detail=dict(_POLICY_BLOCK_DETAIL)) -@pytest.mark.asyncio @pytest.mark.parametrize( "produce,expected_error", [ - (_rate_limit_failure_once_released, {"message": "rate limited", "code": "429"}), - (_guardrail_block_body_once_released, {"message": "blocked", "type": "guardrail_error", "code": "400"}), + (_slow_rate_limit_failure, {"message": "rate limited", "code": "429"}), + (_slow_guardrail_block_body, {"message": "blocked", "type": "guardrail_error", "code": "400"}), ( - _dict_detail_failure_once_released, + _slow_dict_detail_failure, {"message": "blocked by policy", "code": "400", "provider_specific_fields": dict(_POLICY_BLOCK_DETAIL)}, ), ], ) -async def test_ttft_keepalive_ends_a_late_responses_failure_with_an_error_event_and_no_done( - produce: Callable[[asyncio.Event], Awaitable[JSONResponse]], expected_error: Mapping[str, object] +def test_ttft_keepalive_ends_a_late_responses_failure_with_an_error_event_and_no_done( + produce: Callable[[], Awaitable[Response]], expected_error: Mapping[str, object] ) -> None: - released: Final = asyncio.Event() - response: Final = await open_sse_before_first_byte( - produce(released), ping_interval_seconds=0.05, late_framing=RESPONSES_LATE_SSE_FRAMING - ) - assert isinstance(response, StreamingResponse) - frames: Final = aiter(response.body_iterator) - first_frame: Final = await anext(frames) - released.set() - collected: Final = [first_frame, *[frame async for frame in frames]] + 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") - event: Final = json.loads(data_line.removeprefix("data: ")) - assert event_line == "event: error" - assert event["type"] == "error" - assert {key: event["error"].get(key) for key in expected_error} == expected_error + assert collected[0] == TTFT_PING + assert b"data: [DONE]\n\n" not in collected + event_line, data_line = collected[-1].decode().strip().split("\n") + event: Final = json.loads(data_line.removeprefix("data: ")) + assert event_line == "event: error" + assert event["type"] == "error" + assert {key: event["error"].get(key) for key in expected_error} == expected_error + + loop: Final = AutojumpClockLoop() + try: + loop.run_until_complete(collect_and_check()) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() @pytest.mark.asyncio