Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuneng-berri merged 2 commits into from
Jul 18, 2026
Merged

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

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a67d612
test(proxy): make streaming-cancel mocks awaitable for the disconnect…
yuneng-berri Jul 18, 2026
eab54b2
Merge remote-tracking branch 'origin/litellm_internal_staging' into l…
yuneng-berri Jul 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions tests/test_litellm/proxy/test_budget_reservation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2279,14 +2279,16 @@ 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(),
proxy_logging_obj=streaming_logging_obj,
serialize_chunk=lambda chunk: chunk,
serialize_error=lambda exc: str(exc),
)
return generator, streaming_logging_obj


@pytest.mark.asyncio
Expand All @@ -2305,7 +2307,7 @@ async def cancel_before_chunk(user_api_key_dict, response, request_data):
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:
Expand All @@ -2318,6 +2320,7 @@ async def cancel_before_chunk(user_api_key_dict, response, request_data):
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
Expand All @@ -2336,7 +2339,7 @@ async def cancel_after_chunk(user_api_key_dict, response, request_data):
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:
Expand All @@ -2348,6 +2351,7 @@ async def cancel_after_chunk(user_api_key_dict, response, request_data):
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
Expand Down Expand Up @@ -2382,6 +2386,7 @@ async def one_chunk(user_api_key_dict, response, request_data):

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(
Expand Down Expand Up @@ -2411,6 +2416,7 @@ async def one_chunk(user_api_key_dict, response, request_data):
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
Expand All @@ -2427,7 +2433,7 @@ async def two_chunks(user_api_key_dict, response, request_data):
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.
Expand All @@ -2440,6 +2446,7 @@ async def two_chunks(user_api_key_dict, response, request_data):
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
Expand Down
Loading