mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
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.
This commit is contained in:
parent
7ac5002f8d
commit
836a9fe0cf
2 changed files with 8 additions and 15 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue