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

[yuneng-berri]

Relevant issues

Summary

Failure Path (Before Fix)

test_router_caching_ttl intermittently fails with AttributeError: 'NoneType' object has no attribute 'kwargs' when reading mock_client.call_args. The mocked async_increment_cache_pipeline is invoked from Router.deployment_callback_on_success, registered as an async success callback. Those callbacks are enqueued to GLOBAL_LOGGING_WORKER and processed on a background task, so they may not have executed when the test asserts on the mock.

LoggingWorker.flush() also had a latent race: while not self._queue.empty(): await self._queue.join() skipped the wait entirely when the worker had already dequeued a task but not yet called task_done().

Fix

• Drain GLOBAL_LOGGING_WORKER in the test after router.acompletion(...) before asserting on the mock.
• Fix LoggingWorker.flush() to unconditionally await self._queue.join(). asyncio.Queue.join() tracks _unfinished_tasks (incremented by put, decremented by task_done), so it already handles the in-flight case — the empty() guard was wrong.

Testing

• Ran test_router_caching_ttl 5x locally — passes consistently.

Type

🐛 Bug Fix
✅ Test

[veria-ai]

Low: Test flakiness fix with no security impact

This PR fixes a flaky test by improving the flush() method in the logging worker to use queue.join() instead of polling queue.empty(), and adds a flush call in the test. Both changes are limited to test reliability and internal queue management with no security implications.

---

Status: 0 open
Risk: 1/10

Posted by Veria AI · 2026-04-23T22:07:15.960Z

[greptile-apps]

Greptile Summary

This PR fixes intermittent AttributeError: 'NoneType' object has no attribute 'kwargs' in test_router_caching_ttl by draining GLOBAL_LOGGING_WORKER before mock assertions, and simultaneously corrects a race window in flush() that the previous review identified: the old while not self._queue.empty() guard could skip join() if a task was dequeued but task_done() hadn't been called yet — the new implementation calls await self._queue.join() directly, which correctly blocks on the internal _unfinished_tasks counter regardless of queue emptiness.

Confidence Score: 5/5

Safe to merge — the flush() race condition flagged in the previous review is fully addressed, and the test fix is accurate and non-weakening.

Both changes are correct: flush() now uses queue.join() directly which properly waits on the _unfinished_tasks counter rather than the misleading empty() check, and the test drain call correctly eliminates the callback/assertion race. No P0/P1 findings remain.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/litellm_core_utils/logging_worker.py	Fixes flush() race window by replacing the empty()-guarded while loop with a direct await self._queue.join(), correctly handling the dequeued-but-not-done state.
tests/local_testing/test_tpm_rpm_routing_v2.py	Adds await GLOBAL_LOGGING_WORKER.flush() after router.acompletion() to drain the async callback queue before asserting on the mock, eliminating the flakiness race.

Sequence Diagram

sequenceDiagram
    participant Test
    participant Router
    participant GLOBAL_LOGGING_WORKER
    participant asyncio.Queue
    participant Callback as deployment_callback_on_success

    Test->>Router: await acompletion(...)
    Router->>GLOBAL_LOGGING_WORKER: enqueue(callback coroutine)
    GLOBAL_LOGGING_WORKER->>asyncio.Queue: put(task) [_unfinished_tasks++]
    Router-->>Test: response returned
    
    Note over Test,asyncio.Queue: Before fix: test asserts here — callback may not have run yet
    
    Test->>GLOBAL_LOGGING_WORKER: await flush()
    GLOBAL_LOGGING_WORKER->>asyncio.Queue: await join() [blocks until _unfinished_tasks == 0]
    asyncio.Queue->>Callback: worker dequeues and runs callback
    Callback->>asyncio.Queue: task_done() [_unfinished_tasks--]
    asyncio.Queue-->>GLOBAL_LOGGING_WORKER: join() returns
    GLOBAL_LOGGING_WORKER-->>Test: flush() returns
    
    Test->>Test: assert mock_client.call_args ✓

Loading

_{Reviews (2): Last reviewed commit: "fix: make LoggingWorker.flush() wait for..." | Re-trigger Greptile}

[greptile-apps]

flush() has a race window that can still allow early return

The flush() implementation checks self._queue.empty() before calling join(). If the worker loop has already dequeued the item (queue becomes empty) but hasn't yet called task_done(), flush() will skip join() entirely and return before the callback finishes. In that scenario the mock assertion immediately below would still fail.

asyncio.Queue.join() already handles an empty queue correctly (returns immediately when _unfinished_tasks == 0), so the guard is both redundant in the normal path and wrong in the dequeued-but-not-done path. A safer flush() would just be:

async def flush(self) -> None:
    if self._queue is None:
        return
    await self._queue.join()  # waits for all in-flight task_done() calls too

Since flush() is in a shared module and other callers could hit the same edge case, fixing it there is lower risk than relying on scheduling luck in every call site.

[codecov]

Codecov Report

❌ Patch coverage is 94.41088% with 37 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/router_strategy/adaptive_router/hooks.py	84.67%	19 Missing ⚠️
...router_strategy/adaptive_router/adaptive_router.py	96.22%	6 Missing ⚠️
litellm/router_strategy/adaptive_router/signals.py	97.24%	4 Missing ⚠️
...lm/router_strategy/adaptive_router/update_queue.py	93.54%	4 Missing ⚠️
litellm/litellm_core_utils/core_helpers.py	92.85%	2 Missing ⚠️
litellm/router_strategy/adaptive_router/bandit.py	96.61%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!