From 836a9fe0cfabc24fb920c9c06542aaedf230b358 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Mon, 7 Sep 2026 02:21:56 -0700 Subject: [PATCH] test: drop mutable recorder locals and routine comments from the new tests An AsyncMock await_count and an asyncio.Event say the same thing as a list and a dict that the test mutates. --- .../proxy/test_common_request_processing.py | 15 +++++++-------- tests/test_litellm/test_router.py | 8 +------- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 9fcc680b98e..6acd9d7258e 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1816,15 +1816,14 @@ class TestCommonRequestProcessingHelpers: headers are still uncommitted, so ``refresh_headers`` is consulted after the first chunk is buffered and its result wins. """ - refreshed_at: Final[list[str]] = [] async def mock_generator(): yield 'data: {"content": "data"}\n\n' yield "data: [DONE]\n\n" - async def refresh_headers(): - refreshed_at.append("called") - return {"x-litellm-model-id": "fallback-deployment", "llm_provider-x-request-id": "req-FALLBACK"} + refresh_headers: Final = AsyncMock( + return_value={"x-litellm-model-id": "fallback-deployment", "llm_provider-x-request-id": "req-FALLBACK"} + ) response = await create_response( mock_generator(), @@ -1833,7 +1832,7 @@ class TestCommonRequestProcessingHelpers: refresh_headers=refresh_headers, ) assert isinstance(response, StreamingResponse) - assert refreshed_at == ["called"] + assert refresh_headers.await_count == 1 assert response.headers["x-litellm-model-id"] == "fallback-deployment" assert response.headers["llm_provider-x-request-id"] == "req-FALLBACK" # the buffering headers are still applied on top of the refreshed set @@ -1847,15 +1846,15 @@ class TestCommonRequestProcessingHelpers: being produced, so a refresh consulted any earlier still describes the attempt that failed and the headers go out wrong. """ - produced = {"first_chunk": False} + first_chunk_produced: Final = asyncio.Event() async def mock_generator(): - produced["first_chunk"] = True + first_chunk_produced.set() yield 'data: {"content": "data"}\n\n' yield "data: [DONE]\n\n" async def refresh_headers(): - served = "fallback-deployment" if produced["first_chunk"] else "failed-deployment" + served = "fallback-deployment" if first_chunk_produced.is_set() else "failed-deployment" return {"x-litellm-model-id": served} response = await create_response( diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 70ebdc4766c..3bc7155952c 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2524,12 +2524,7 @@ def test_adopt_fallback_response_headers_replaces_rather_than_merges(): def test_adopt_fallback_response_headers_survives_a_collected_wrapper(): - """LIT-6767: the fallback generator holds only a weak reference to its wrapper. - - A client that disconnects mid-stream can drop the wrapper while the generator is - still draining, and the adoption call has to keep working with nothing to adopt into - so the deployment slot is still released. - """ + """LIT-6767: adoption still returns the fallback's params once the wrapper is gone.""" import weakref from unittest.mock import MagicMock @@ -2557,7 +2552,6 @@ def test_adopt_fallback_response_headers_survives_a_collected_wrapper(): dead_ref: Final = weakref.ref(wrapper) del wrapper assert dead_ref() is None - # no wrapper left to repoint, and the caller still needs the params for the chunks assert Router._adopt_fallback_response_headers(dead_ref, fallback) == prepared