Merge pull request #26355 from BerriAI/litellm_fixFlakyTpmRoutingTest

[Fix] Tests - drain logging worker in test_router_caching_ttl to fix flakiness

litellm/litellm_core_utils/logging_worker.py

@@ -370,11 +370,17 @@ class LoggingWorker:
         self._running_tasks.clear()
 
     async def flush(self) -> None:
-        """Flush the logging queue."""
+        """Flush the logging queue.
+
+        Waits until every enqueued task has completed. ``queue.join()`` blocks
+        on the queue's unfinished-task counter (decremented by ``task_done()``),
+        so it correctly handles items that have been dequeued but whose
+        callback hasn't finished yet — ``queue.empty()`` would return True in
+        that window and cause us to skip the wait.
+        """
         if self._queue is None:
             return
-        while not self._queue.empty():
-            await self._queue.join()
+        await self._queue.join()
 
     async def clear_queue(self):
         """

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