scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

test(streaming): restore success callbacks and await dispatch deterministically

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
milan 2026-07-28 23:16:37 +00:00
parent ef26590f72
commit 09e1fb5ea2
1 changed files with 15 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
View file

@ -1553,6 +1553,8 @@ async def test_openrouter_streaming_usage_only_chunk_without_stream_options():
pass
mock_callback = MockCallback()
previous_success_callback = litellm.success_callback
previous_async_success_callback = litellm._async_success_callback
litellm.success_callback = [mock_callback]
litellm._async_success_callback = [mock_callback]
@ -1581,9 +1583,19 @@ async def test_openrouter_streaming_usage_only_chunk_without_stream_options():
logging_obj=stream_logging_obj,
)
with patch.object(mock_callback, "async_log_success_event") as mock_success_event:
collected_chunks = [chunk async for chunk in response]
await asyncio.sleep(1)
success_logged = asyncio.Event()
try:
with patch.object(
mock_callback,
"async_log_success_event",
new_callable=AsyncMock,
side_effect=lambda *args, **kwargs: success_logged.set(),
) as mock_success_event:
collected_chunks = [chunk async for chunk in response]
await asyncio.wait_for(success_logged.wait(), timeout=30)
finally:
litellm.success_callback = previous_success_callback
litellm._async_success_callback = previous_async_success_callback
assert all(getattr(chunk, "usage", None) is None for chunk in collected_chunks)