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 109

test(rate-limiting): use monkeypatch.setattr instead of manual litellm.callbacks restore

Browse source 
The base added a new test-quality gate (TQ005) flagging direct
litellm.<attr> = ... module-global mutation, even when restored via
try/finally. Switches both disconnect-hook tests to monkeypatch.setattr,
which the gate doesn't flag and which pytest reverts automatically at
teardown, removing the manual restore entirely.
This commit is contained in:
Deepanshu 2026-08-20 14:11:13 -04:00
parent 0c07a26c8f
commit e7d5b51a39
1 changed files with 27 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -6000,7 +6000,7 @@ class TestStreamingClientDisconnectBilling:
assert usage.prompt_tokens_details.cached_tokens == 7
@pytest.mark.asyncio
async def test_disconnect_without_billable_chunks_releases_callback_state(self):
async def test_disconnect_without_billable_chunks_releases_callback_state(self, monkeypatch):
"""
A callback that reserves per-request state outside of the success/failure
logging callbacks (e.g. a concurrency slot admitted before the first
@ -6014,27 +6014,23 @@ class TestStreamingClientDisconnectBilling:
response = await self._start_partial_stream()
empty_response = types.SimpleNamespace(chunks=[], messages=None)
recorder = _RecordingDisconnectHookLogger()
original_callbacks = litellm.callbacks
litellm.callbacks = [recorder]
try:
await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup(
request=None,
request_data={"litellm_logging_obj": response.logging_obj},
response=empty_response,
stream_completed=False,
client_disconnected=True,
user_api_key_dict=MagicMock(),
proxy_logging_obj=types.SimpleNamespace(
_arelease_max_parallel_requests_on_disconnect=AsyncMock(),
),
)
finally:
litellm.callbacks = original_callbacks
monkeypatch.setattr(litellm, "callbacks", [recorder])
await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup(
request=None,
request_data={"litellm_logging_obj": response.logging_obj},
response=empty_response,
stream_completed=False,
client_disconnected=True,
user_api_key_dict=MagicMock(),
proxy_logging_obj=types.SimpleNamespace(
_arelease_max_parallel_requests_on_disconnect=AsyncMock(),
),
)
assert recorder.disconnect_hook_calls == 1
@pytest.mark.asyncio
async def test_disconnect_billing_skips_callback_disconnect_hook(self):
async def test_disconnect_billing_skips_callback_disconnect_hook(self, monkeypatch):
"""
When a disconnect-time success event already fired (partial billing
dispatched it), that event's own async_log_success_event already ran
@ -6045,23 +6041,19 @@ class TestStreamingClientDisconnectBilling:
import types
recorder = _RecordingDisconnectHookLogger()
original_callbacks = litellm.callbacks
litellm.callbacks = [recorder]
try:
response = await self._start_partial_stream()
await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup(
request=None,
request_data={"litellm_logging_obj": response.logging_obj},
response=response,
stream_completed=False,
client_disconnected=True,
user_api_key_dict=MagicMock(),
proxy_logging_obj=types.SimpleNamespace(
_arelease_max_parallel_requests_on_disconnect=AsyncMock(),
),
)
finally:
litellm.callbacks = original_callbacks
monkeypatch.setattr(litellm, "callbacks", [recorder])
response = await self._start_partial_stream()
await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup(
request=None,
request_data={"litellm_logging_obj": response.logging_obj},
response=response,
stream_completed=False,
client_disconnected=True,
user_api_key_dict=MagicMock(),
proxy_logging_obj=types.SimpleNamespace(
_arelease_max_parallel_requests_on_disconnect=AsyncMock(),
),
)
assert recorder.disconnect_hook_calls == 0