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

[yuneng-berri]

Relevant issues

Repairs the proxy-infra CI breakage introduced by #33736, which landed while its proxy-infra check was still failing and has been breaking litellm_internal_staging and every PR on top of it since

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

Screenshots / Proof of Fix

This PR changes only a unit test file, so the observable surface is the proxy-infra check itself; no production code changes, hence no live-proxy repro applies

Before, without this fix: the proxy-infra job fails with the same four tests on staging head ca720c4 (first link) and again on #33801 whose CI merge ref sat on staging tip 577dd3b (second link)

• https://github.com/BerriAI/litellm/actions/runs/29622577395
• https://github.com/BerriAI/litellm/actions/runs/29624783051/job/88026890492

FAILED tests/test_litellm/proxy/test_budget_reservation.py::test_streaming_cancel_after_chunk_keeps_reservation - TypeError: object MagicMock can't be used in 'await' expression
FAILED tests/test_litellm/proxy/test_budget_reservation.py::test_streaming_cancel_before_any_chunk_reconciles_to_input_cost - TypeError: object MagicMock can't be used in 'await' expression
FAILED tests/test_litellm/proxy/test_budget_reservation.py::test_streaming_cancel_in_slow_path_before_yield_refunds - TypeError: object MagicMock can't be used in 'await' expression
FAILED tests/test_litellm/proxy/test_budget_reservation.py::test_streaming_disconnect_after_consuming_chunk_keeps_reservation - TypeError: object MagicMock can't be used in 'await' expression

After, at a67d612 (PR head eab54b2): the proxy-infra check on this PR runs the same suite and passes in 8m1s

• https://github.com/BerriAI/litellm/actions/runs/29625659037/job/88029418267

The osv-scan failure on this PR is the same pre-existing repo-wide failure visible on every current PR (mcp 1.26.0 advisories in uv.lock, fixed releases available upstream); it is unrelated to this change and tracked separately

Type

✅ Test

Changes

#33736 made the shielded streaming cleanup in common_request_processing.py await proxy_logging_obj._arelease_max_parallel_requests_on_disconnect when a client disconnect is recorded and no disconnect-time success event owns the slot release. The four streaming cancel and disconnect tests in tests/test_litellm/proxy/test_budget_reservation.py drive async_streaming_data_generator with a bare MagicMock as proxy_logging_obj, and awaiting a MagicMock call result raises TypeError, so all four fail on every cancel and disconnect path

The fix gives those mocks an AsyncMock for _arelease_max_parallel_requests_on_disconnect, both in the shared _drive_streaming_cancel helper and in the slow-path test that builds its mock inline. Each of the four tests now also asserts the release is awaited exactly once, pinning the single-owner slot release contract #33736 introduced without test coverage. Verified by mutation: removing the release call from _finalize_streaming_generator_cleanup makes these tests fail again, and restoring it turns the whole file green (51 passed)

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

[greptile-apps]

Greptile Summary

This PR fixes a CI breakage introduced by #33736 that caused four async tests to fail with TypeError: object MagicMock can't be used in 'await' expression. The fix adds AsyncMock for _arelease_max_parallel_requests_on_disconnect on the shared _drive_streaming_cancel helper mock and on the inline mock in the slow-path test, so the await added by #33736 can succeed.

• _drive_streaming_cancel is refactored to return both the generator and the streaming_logging_obj, allowing each test to assert assert_awaited_once() on the release mock.
• All four affected tests now also assert that _arelease_max_parallel_requests_on_disconnect is awaited exactly once, adding test-pinned coverage for the single-owner slot-release contract introduced by fix(proxy): bill partial streamed spend when the client disconnects mid-stream #33736.
• No production code is touched; all changes are confined to the test file.

Confidence Score: 5/5

Safe to merge — only the test file is changed, and the changes fix, not weaken, existing test coverage.

The patch is test-only. It makes previously-crashing tests runnable by providing an awaitable mock and adds assert_awaited_once() assertions that strengthen rather than weaken coverage. The _drive_streaming_cancel refactor is minimal and mechanical. There are no changes to production code paths, and the PR description includes clear before/after CI evidence.

No files require special attention.

Important Files Changed

Filename	Overview
tests/test_litellm/proxy/test_budget_reservation.py	Fixes four broken streaming-cancel tests by adding AsyncMock for _arelease_max_parallel_requests_on_disconnect and adds assertions that the slot is released exactly once per cancel/disconnect path.

_{Reviews (1): Last reviewed commit: "Merge remote-tracking branch 'origin/lit..." | Re-trigger Greptile}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_/pr-failing-tests-dc7207 (eab54b2) with litellm_internal_staging (f3d2015)¹}

Footnotes

1. No successful run was found on litellm_internal_staging (967d934) during the generation of this report, so f3d2015 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩