From 394818e45be707c6e2f574c4e2ee0772c0548f76 Mon Sep 17 00:00:00 2001 From: yryzhan-vitech Date: Mon, 27 Jul 2026 22:26:58 +0200 Subject: [PATCH] test(proxy): pin that the scheduler stops before the shutdown drain flushes Nothing asserted the ordering between stopping the scheduler and flushing, so moving the stop after the flush loop still passed while letting a scheduled job race the final drain. --- .../proxy/proxy_server/test_lifecycle.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index e34e5e63e50..6d89cc55303 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -865,3 +865,25 @@ async def test_drain_spend_buffers_propagates_cancellation(monkeypatch): await ps._drain_spend_buffers_on_shutdown() assert deps.tag_spend.await_count == 0 + + +@pytest.mark.asyncio +async def test_drain_spend_buffers_stops_scheduler_before_flushing(monkeypatch): + order: List[str] = [] + + scheduler = MagicMock() + scheduler.shutdown.side_effect = lambda **kwargs: order.append("scheduler_stopped") + + async def _flush(*args, **kwargs): + order.append("flush") + + _patch_drain_deps( + monkeypatch, + scheduler=scheduler, + update_spend=AsyncMock(side_effect=_flush), + update_daily_tag_spend=AsyncMock(side_effect=_flush), + ) + + await ps._drain_spend_buffers_on_shutdown() + + assert order == ["scheduler_stopped", "flush", "flush"]