litellm/tests/_wait_helpers.py
ryan-crabbe-berri 487356733c
test: replace blind sleeps with deadline waits in callback and caching tests (#37660)
* 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.
2026-08-20 18:48:43 +00:00

46 lines
1.3 KiB
Python

"""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)