scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

test(proxy): make streaming-cancel mocks awaitable for the disconnect slot release (#33802)

Browse source 
PR #33736 made the shielded streaming cleanup await
proxy_logging_obj._arelease_max_parallel_requests_on_disconnect on the
client-disconnect path. The four streaming cancel and disconnect tests in
test_budget_reservation.py drive the generator with a bare MagicMock as
proxy_logging_obj, so the cleanup crashed with TypeError: object MagicMock
can't be used in 'await' expression, breaking proxy-infra CI on every PR

Give the mocks an AsyncMock for the release method and assert it is awaited
exactly once on each disconnect path, pinning the single-owner slot release
contract that PR #33736 introduced without test coverage
This commit is contained in:
yuneng-jiang 2026-07-17 18:50:20 -07:00 • committed by GitHub
parent 967d934484
commit a4c9571181
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 11 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -2279,7 +2279,8 @@ async def _reserve_for_stream(counter_cache, key_cache, proxy_logging_obj, token
def _drive_streaming_cancel(valid_token, iterator_hook):
streaming_logging_obj = MagicMock()
streaming_logging_obj.async_post_call_streaming_iterator_hook = iterator_hook
return ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
streaming_logging_obj._arelease_max_parallel_requests_on_disconnect = AsyncMock()
generator = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
response=MagicMock(),
user_api_key_dict=valid_token,
request_data=_request_body(),
@ -2287,6 +2288,7 @@ def _drive_streaming_cancel(valid_token, iterator_hook):
serialize_chunk=lambda chunk: chunk,
serialize_error=lambda exc: str(exc),
)
return generator, streaming_logging_obj
@pytest.mark.asyncio
@ -2305,7 +2307,7 @@ async def test_streaming_cancel_before_any_chunk_reconciles_to_input_cost(
yield "" # make this an async generator
raise asyncio.CancelledError()
generator = _drive_streaming_cancel(valid_token, cancel_before_chunk)
generator, streaming_logging_obj = _drive_streaming_cancel(valid_token, cancel_before_chunk)
received = []
with pytest.raises(asyncio.CancelledError):
async for chunk in generator:
@ -2318,6 +2320,7 @@ async def test_streaming_cancel_before_any_chunk_reconciles_to_input_cost(
key="spend:key:key-cancel-no-chunk"
) == pytest.approx(0.5)
assert reservation["finalized"] is True
streaming_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once()
@pytest.mark.asyncio
@ -2336,7 +2339,7 @@ async def test_streaming_cancel_after_chunk_keeps_reservation(
yield "data: chunk\n\n"
raise asyncio.CancelledError()
generator = _drive_streaming_cancel(valid_token, cancel_after_chunk)
generator, streaming_logging_obj = _drive_streaming_cancel(valid_token, cancel_after_chunk)
received = []
with pytest.raises(asyncio.CancelledError):
async for chunk in generator:
@ -2348,6 +2351,7 @@ async def test_streaming_cancel_after_chunk_keeps_reservation(
key="spend:key:key-cancel-after-chunk"
) == pytest.approx(2.0)
assert reservation.get("finalized") is not True
streaming_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once()
@pytest.mark.asyncio
@ -2382,6 +2386,7 @@ async def test_streaming_cancel_in_slow_path_before_yield_refunds(spend_counter_
streaming_logging_obj = MagicMock()
streaming_logging_obj.async_post_call_streaming_iterator_hook = one_chunk
streaming_logging_obj._arelease_max_parallel_requests_on_disconnect = AsyncMock()
# On the slow path the per-chunk hook is awaited before the chunk is yielded
# to the client; cancel there. Nothing has reached the client yet.
streaming_logging_obj.async_post_call_streaming_hook = AsyncMock(
@ -2411,6 +2416,7 @@ async def test_streaming_cancel_in_slow_path_before_yield_refunds(spend_counter_
key="spend:key:key-cancel-slowpath"
) == pytest.approx(0.5)
assert reservation["finalized"] is True
streaming_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once()
@pytest.mark.asyncio
@ -2427,7 +2433,7 @@ async def test_streaming_disconnect_after_consuming_chunk_keeps_reservation(
yield "data: a\n\n"
yield "data: b\n\n"
generator = _drive_streaming_cancel(valid_token, two_chunks)
generator, streaming_logging_obj = _drive_streaming_cancel(valid_token, two_chunks)
# Client consumes one chunk, then disconnects. aclose() raises GeneratorExit
# at the suspended yield, after the chunk already reached the client.
@ -2440,6 +2446,7 @@ async def test_streaming_disconnect_after_consuming_chunk_keeps_reservation(
key="spend:key:key-disconnect-after-chunk"
) == pytest.approx(2.0)
assert reservation.get("finalized") is not True
streaming_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once()
@pytest.mark.asyncio