fix: drain logging worker in test_router_caching_ttl to remove flake

The mocked async_increment_cache_pipeline is invoked from Router's
deployment_callback_on_success, registered as an async success callback.
Those callbacks are enqueued to GLOBAL_LOGGING_WORKER and run on a
background task, so the mock may not have been called yet when the test
asserts on it. Flush the worker before asserting.
This commit is contained in:
Yuneng Jiang 2026-04-23 14:48:02 -07:00
parent e9e86ed956
commit c4ea0e93c8
No known key found for this signature in database

View file

@ -547,6 +547,8 @@ async def test_router_caching_ttl():
assert router.cache.redis_cache is not None
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
increment_cache_kwargs = {}
with patch.object(
router.cache,
@ -555,6 +557,10 @@ async def test_router_caching_ttl():
) as mock_client:
await router.acompletion(model=model, messages=messages)
# Async success callbacks are dispatched to GLOBAL_LOGGING_WORKER's
# background queue; drain it before asserting the mock was invoked.
await GLOBAL_LOGGING_WORKER.flush()
# mock_client.assert_called_once()
print(f"mock_client.call_args.kwargs: {mock_client.call_args.kwargs}")
print(f"mock_client.call_args.args: {mock_client.call_args.args}")