From a4c9571181f12e5d7d0dc6f2e69a21a2ad89aba3 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng@berri.ai>
Date: Fri, 17 Jul 2026 18:50:20 -0700
Subject: [PATCH] test(proxy): make streaming-cancel mocks awaitable for the
 disconnect slot release (#33802)

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
---
 .../test_litellm/proxy/test_budget_reservation.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py
index 0b304f2fec7..1db76aed61d 100644
--- a/tests/test_litellm/proxy/test_budget_reservation.py
+++ b/tests/test_litellm/proxy/test_budget_reservation.py
@@ -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