mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
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.
This commit is contained in:
parent
7c84adeec2
commit
394818e45b
1 changed files with 22 additions and 0 deletions
|
|
@ -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"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue