From c4ea0e93c80772d1615e36eb1052b2179465885f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 23 Apr 2026 14:48:02 -0700 Subject: [PATCH] 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. --- tests/local_testing/test_tpm_rpm_routing_v2.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/local_testing/test_tpm_rpm_routing_v2.py b/tests/local_testing/test_tpm_rpm_routing_v2.py index 9de5625c63a..211af566424 100644 --- a/tests/local_testing/test_tpm_rpm_routing_v2.py +++ b/tests/local_testing/test_tpm_rpm_routing_v2.py @@ -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}")