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: replace blind sleeps with deadline waits in callback and caching tests (#37660)

Browse source 
* test: replace blind sleeps with deadline waits in callback and caching tests

tests/local_testing/test_custom_callback_input.py slept a fixed 1-3s after
every call and then asserted the callback handler recorded no errors. Because
the handler only appends to `states` when a callback actually fires, an assert
of `len(errors) == 0` passes just as happily when nothing fired at all, so the
sleep was buying flakiness in exchange for a vacuous check. The async tests
were worse: `time.sleep` blocks the event loop, so the success/failure tasks
scheduled on it could not run before the assertion.

Adds tests/_wait_helpers.py with `wait_until` / `await_until`, which poll a
predicate against a deadline, and converts all 17 sites to wait on the thing
the test actually cares about (the terminal state landing in `states`, or the
patched log hook being called). The waits assert the callback fired, so these
tests now fail on a dropped callback instead of passing silently.

The three sleeps in test_caching_handler.py sat between `sync_set_cache` and
`_sync_get_cache`, both fully synchronous against a local in-memory cache, so
they are just deleted.

* fix(test): wait on the priming call's own logging in the cache-hit test

The 3s sleep in test_logging_async_cache_hit_sync_call was not waiting for the
cache write, which lands before the stream iterator is exhausted. It was
waiting for the priming call's success callback to drain, so the handler
installed right after it only ever sees the second, cache-hit call. Waiting on
a populated cache_dict let the priming call's still-pending log_success_event
reach the new mock, and the test then read cache_hit off the wrong payload.

Waits on the priming handler's own sync_success state instead.
This commit is contained in:
ryan-crabbe-berri 2026-08-20 11:48:43 -07:00 • committed by GitHub
parent d542c82f0e
commit 487356733c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 103 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
tests/_wait_helpers.py Normal file
View file

@ -0,0 +1,46 @@
"""Deadline-based waits for tests, so nothing has to guess how long a background callback takes."""
import asyncio
import time
from collections.abc import Callable
from typing import Final
DEFAULT_TIMEOUT_S: Final[float] = 10.0
DEFAULT_INTERVAL_S: Final[float] = 0.02
def _fail(timeout_s: float, message: str) -> None:
raise AssertionError(f"condition not met within {timeout_s}s: {message}")
def wait_until(
predicate: Callable[[], bool],
*,
message: str,
timeout_s: float = DEFAULT_TIMEOUT_S,
interval_s: float = DEFAULT_INTERVAL_S,
) -> None:
deadline: Final = time.monotonic() + timeout_s
while time.monotonic() < deadline:
if predicate():
return
time.sleep(interval_s) # sleep-ok: bounded poll interval, not a blind settle
if not predicate():
_fail(timeout_s, message)
async def await_until(
predicate: Callable[[], bool],
*,
message: str,
timeout_s: float = DEFAULT_TIMEOUT_S,
interval_s: float = DEFAULT_INTERVAL_S,
) -> None:
"""Yields to the event loop between polls, so callbacks scheduled as tasks get a chance to run."""
deadline: Final = time.monotonic() + timeout_s
while time.monotonic() < deadline:
if predicate():
return
await asyncio.sleep(interval_s)
if not predicate():
_fail(timeout_s, message)

4
tests/local_testing/test_caching_handler.py
View file

@ -741,8 +741,6 @@ def test_sync_responses_api_caching():
# Step 1: Cache the responses API response
caching_handler.sync_set_cache(result=responses_api_response, kwargs=kwargs)
time.sleep(0.5)
# Step 2: Retrieve from cache
cached_response = caching_handler._sync_get_cache(
model=original_model,
@ -875,7 +873,6 @@ def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits():
}
caching_handler.sync_set_cache(result=responses_api_response, kwargs=kwargs)
time.sleep(0.2)
cached_response = caching_handler._sync_get_cache(
model=original_model,
@ -920,7 +917,6 @@ def test_sync_get_cache_defers_streaming_completion_hit_callbacks():
}
caching_handler.sync_set_cache(result=chat_completion_response, kwargs=kwargs)
time.sleep(0.2)
cached_response = caching_handler._sync_get_cache(
model=original_model,

79
tests/local_testing/test_custom_callback_input.py
View file

@ -4,7 +4,6 @@ import asyncio
import inspect
import os
import sys
import time
import traceback
from litellm._uuid import uuid
from datetime import datetime
@ -20,6 +19,7 @@ import litellm
from litellm import Cache, completion, embedding
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import LiteLLMCommonStrings
from tests._wait_helpers import await_until, wait_until
# Test Scenarios (test across completion, streaming, embedding)
## 1: Pre-API-Call
@ -389,7 +389,10 @@ def test_chat_openai_stream():
continue
except Exception:
pass
time.sleep(1)
wait_until(
lambda: "sync_failure" in customHandler.states,
message=f"no sync_failure callback, states={customHandler.states}",
)
print(f"customHandler.errors: {customHandler.errors}")
assert len(customHandler.errors) == 0
litellm.callbacks = []
@ -430,10 +433,12 @@ async def test_async_chat_openai_stream():
)
async for chunk in response:
continue
await asyncio.sleep(1)
except Exception:
pass
time.sleep(1)
await await_until(
lambda: "async_failure" in customHandler.states,
message=f"no async_failure callback, states={customHandler.states}",
)
print(f"customHandler.errors: {customHandler.errors}")
assert len(customHandler.errors) == 0
litellm.callbacks = []
@ -473,7 +478,10 @@ def test_chat_azure_stream():
continue
except Exception:
pass
time.sleep(1)
wait_until(
lambda: "sync_failure" in customHandler.states,
message=f"no sync_failure callback, states={customHandler.states}",
)
print(f"customHandler.errors: {customHandler.errors}")
assert len(customHandler.errors) == 0
litellm.callbacks = []
@ -590,7 +598,10 @@ async def test_async_chat_sagemaker_stream():
continue
except Exception:
pass
time.sleep(1)
await await_until(
lambda: "async_failure" in customHandler.states,
message=f"no async_failure callback, states={customHandler.states}",
)
print(f"customHandler.errors: {customHandler.errors}")
assert len(customHandler.errors) == 0
litellm.callbacks = []
@ -711,10 +722,12 @@ async def test_async_text_completion_bedrock():
async for chunk in response:
continue
await asyncio.sleep(1)
except Exception:
pass
time.sleep(1)
await await_until(
lambda: "async_failure" in customHandler.states,
message=f"no async_failure callback, states={customHandler.states}",
)
print(f"customHandler.errors: {customHandler.errors}")
assert len(customHandler.errors) == 0
litellm.callbacks = []
@ -754,10 +767,12 @@ async def test_async_text_completion_openai_stream():
async for chunk in response:
continue
await asyncio.sleep(1)
except Exception:
pass
time.sleep(1)
await await_until(
lambda: "async_failure" in customHandler.states,
message=f"no async_failure callback, states={customHandler.states}",
)
print(f"customHandler.errors: {customHandler.errors}")
assert len(customHandler.errors) == 0
litellm.callbacks = []
@ -816,7 +831,10 @@ def test_amazing_sync_embedding():
)
print(f"customHandler_success.errors: {customHandler_success.errors}")
print(f"customHandler_success.states: {customHandler_success.states}")
time.sleep(2)
wait_until(
lambda: len(customHandler_success.states) == 3,
message=f"success states never reached pre/post/success, got {customHandler_success.states}",
)
assert len(customHandler_success.errors) == 0
assert len(customHandler_success.states) == 3 # pre, post, success
# test failure callback
@ -832,7 +850,10 @@ def test_amazing_sync_embedding():
pass
print(f"customHandler_failure.errors: {customHandler_failure.errors}")
print(f"customHandler_failure.states: {customHandler_failure.states}")
time.sleep(2)
wait_until(
lambda: len(customHandler_failure.states) == 3,
message=f"failure states never reached pre/post/failure, got {customHandler_failure.states}",
)
assert len(customHandler_failure.errors) == 1
assert len(customHandler_failure.states) == 3 # pre, post, failure
except Exception as e:
@ -939,7 +960,10 @@ def test_image_generation_openai():
print(f"customHandler_success.errors: {customHandler_success.errors}")
print(f"customHandler_success.states: {customHandler_success.states}")
time.sleep(2)
wait_until(
lambda: len(customHandler_success.states) == 3,
message=f"success states never reached pre/post/success, got {customHandler_success.states}",
)
assert len(customHandler_success.errors) == 0
assert len(customHandler_success.states) == 3 # pre, post, success
# test failure callback
@ -991,7 +1015,10 @@ def test_turn_off_message_logging():
mock_response="Going well!",
)
time.sleep(2)
wait_until(
lambda: "sync_success" in customHandler.states,
message=f"no sync_success callback, states={customHandler.states}",
)
assert len(customHandler.errors) == 0
@ -1033,7 +1060,7 @@ def test_standard_logging_payload(model, turn_off_message_logging):
mock_response="Going well!",
)
time.sleep(2)
wait_until(lambda: mock_client.called, message="log_success_event never fired")
mock_client.assert_called_once()
print(
@ -1147,7 +1174,7 @@ def test_standard_logging_payload_audio(turn_off_message_logging, stream):
for chunk in response:
continue
time.sleep(2)
wait_until(lambda: mock_client.called, message="log_success_event never fired")
mock_client.assert_called()
print(
@ -1247,7 +1274,7 @@ def test_aaastandard_logging_payload_cache_hit():
caching=True,
)
time.sleep(2)
wait_until(lambda: mock_client.called, message="log_success_event never fired")
mock_client.assert_called_once()
assert "standard_logging_object" in mock_client.call_args.kwargs["kwargs"]
@ -1276,6 +1303,9 @@ def test_logging_async_cache_hit_sync_call(turn_off_message_logging):
litellm.cache = Cache()
primingHandler = CompletionCustomHandler()
litellm.callbacks = [primingHandler]
response = litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
@ -1285,7 +1315,10 @@ def test_logging_async_cache_hit_sync_call(turn_off_message_logging):
for chunk in response:
print(chunk)
time.sleep(3)
wait_until(
lambda: "sync_success" in primingHandler.states,
message=f"priming call never finished logging, states={primingHandler.states}",
)
customHandler = CompletionCustomHandler()
litellm.callbacks = [customHandler]
litellm.success_callback = []
@ -1303,7 +1336,7 @@ def test_logging_async_cache_hit_sync_call(turn_off_message_logging):
for chunk in resp:
print(chunk)
time.sleep(2)
wait_until(lambda: mock_client.called, message="log_success_event never fired")
mock_client.assert_called_once()
assert "standard_logging_object" in mock_client.call_args.kwargs["kwargs"]
@ -1387,7 +1420,7 @@ def test_logging_standard_payload_llm_headers(stream):
for chunk in resp:
continue
time.sleep(2)
wait_until(lambda: mock_client.called, message="log_success_event never fired")
mock_client.assert_called()
standard_logging_object: StandardLoggingPayload = mock_client.call_args.kwargs[
@ -1458,7 +1491,7 @@ async def test_standard_logging_payload_stream_usage(sync_mode):
chunks = []
for chunk in resp:
chunks.append(chunk)
time.sleep(2)
wait_until(lambda: mock_client.called, message="log_success_event never fired")
else:
resp = await litellm.acompletion(
model="anthropic/claude-sonnet-4-5-20250929",
@ -1469,7 +1502,9 @@ async def test_standard_logging_payload_stream_usage(sync_mode):
chunks = []
async for chunk in resp:
chunks.append(chunk)
await asyncio.sleep(2)
await await_until(
lambda: mock_client.called, message="async_log_success_event never fired"
)
mock_client.assert_called_once()