From 487356733c7254944e7825bdd7c87e6d8f262e8f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 20 Aug 2026 11:48:43 -0700 Subject: [PATCH] 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. --- tests/_wait_helpers.py | 46 +++++++++++ tests/local_testing/test_caching_handler.py | 4 - .../test_custom_callback_input.py | 79 +++++++++++++------ 3 files changed, 103 insertions(+), 26 deletions(-) create mode 100644 tests/_wait_helpers.py diff --git a/tests/_wait_helpers.py b/tests/_wait_helpers.py new file mode 100644 index 00000000000..f67e623e3ad --- /dev/null +++ b/tests/_wait_helpers.py @@ -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) diff --git a/tests/local_testing/test_caching_handler.py b/tests/local_testing/test_caching_handler.py index 0f4539162a2..b26334e9ee0 100644 --- a/tests/local_testing/test_caching_handler.py +++ b/tests/local_testing/test_caching_handler.py @@ -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, diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py index 6a4ec9206f7..cedb5ea1a97 100644 --- a/tests/local_testing/test_custom_callback_input.py +++ b/tests/local_testing/test_custom_callback_input.py @@ -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()