test(e2e): poll key spend to a deadline in budget reset advances tests

[ryan-crabbe-berri]

TLDR

Problem this solves:

• stage flake: budget reset tests read key spend exactly once
• the ~60s batched spend flush re-lands pre-reset deltas after the reset zeroes the row
• single read races that flush; failed Jul 30 and Aug 2 stage runs with spend ~9.7e-05 vs 3e-06 cap

How it solves it:

• new _poll_key_spend helper polls the spend read to a 90s deadline
• rung 3 (spend zeroed) and rung 4 (roomy window keeps spend) now poll instead of single-read
• a reset that never zeroes the row still fails: spend stays pinned and the poll times out

Relevant issues

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

Screenshots / Proof of Fix

Before (the flake this kills): stage e2e runs 2026-07-30 00:45 UTC and 2026-08-02 00:52 UTC (image at ba480a6) both failed test_key_budget_reset_at_advances_after_window with the identical signature while every neighbouring assertion passed, proving the reset itself worked and only the read raced the flush

E   AssertionError: spend not cleared after reset: 9.7e-05
E   assert ((9.7e-05)) < 3e-06
E    +  where 9.7e-05 = KeyInfo(..., spend=9.7e-05, max_budget=3e-06, budget_reset_at='2026-08-02T01:35:30+00:00', ...).spend

After (at ebdd854): both tests against a live worktree proxy on localhost:4611 backed by a real provider (groq llama-3.3-70b-versatile, x-litellm-response-cost 2.282e-05 per call), with the proxy running the stage cadence that produces the race window (PROXY_BATCH_WRITE_AT=60, budget rescheduler every 15-20s)

LITELLM_PROXY_URL=http://localhost:4611 python -m pytest \
  quota_management/budgets/test_budget_reset_advances_e2e.py::test_key_budget_reset_at_advances_after_window \
  quota_management/budgets/test_budget_reset_advances_e2e.py::test_multi_window_key_resets_each_window_independently -v

quota_management/budgets/test_budget_reset_advances_e2e.py::test_key_budget_reset_at_advances_after_window PASSED [ 50%]
quota_management/budgets/test_budget_reset_advances_e2e.py::test_multi_window_key_resets_each_window_independently PASSED [100%]

======================== 2 passed in 100.88s (0:01:40) =========================

Type

✅ Test

Changes

tests/e2e/quota_management/budgets/test_budget_reset_advances_e2e.py only. Adds _poll_key_spend, which re-reads key_info.spend every 5s until a predicate holds or a 90s deadline expires; 90s covers one full flush-plus-reset cycle, so a late flush delays the pass a few iterations while a reset job that never zeroes the row keeps spend pinned at the driven total and still times out. Rung 3 swaps its single spend < cap read for the poll, rung 4 swaps its single spend >= spend_at_block read (the mirror timing of the same race), and both docstrings say why. No other budget test reads spend as a single post-convergence snapshot: the reseed test already waits out the flush explicitly and the user-budget test already polls

QA runbook

• tests/e2e/quota_management/budgets/test_budget_reset_advances_e2e.py::test_key_budget_reset_at_advances_after_window - after the budget window elapses the reset job advances budget_reset_at and zeroes spend, with the spend read now polled to a 90s deadline
  • Generate a key with max_budget 3e-06 and budget_duration 30s, then note its budget_reset_at from /key/info
  • Send chat completions with that key until one returns the budget_exceeded block
  • Keep retrying a call every 5s; within 150s one succeeds, and /key/info then shows budget_reset_at strictly later than the noted value
  • Re-read /key/info every 5s for up to 90s and expect spend to drop under the cap once the flush and the next reset settle
  • Sanity check: this test makes sense and is not hand-wavey or potentially flaky
• tests/e2e/quota_management/budgets/test_budget_reset_advances_e2e.py::test_multi_window_key_resets_each_window_independently - a tight 30s window resets on its own boundary while the roomy 1m window keeps its accumulated spend
  • Generate a key with two budget windows: 30s at 3e-06 and 1m at 1.0, then drive it to the budget block and note spend from /key/info
  • Keep retrying a call every 5s; within 150s one succeeds and well inside the tight window plus 90s
  • Re-read /key/info every 5s for up to 90s and expect spend to come back at or above the noted value, proving only the tight window reset
  • Sanity check: this test makes sense and is not hand-wavey or potentially flaky

Environment notes: needs a DB-backed proxy whose budget rescheduler runs faster than the default 10 minutes (stage runs 15-20s) and a model named claude-haiku-4-5; locally I pointed that alias at groq since this machine's .env carries no Anthropic key

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

[greptile-apps]

Greptile Summary

This test-only PR adds a reusable deadline-based key-spend poll to tolerate asynchronous spend-write convergence in two budget reset E2E assertions.

• Polls key spend every five seconds for up to 90 seconds.
• Replaces one-shot post-reset spend reads in the single-window and multi-window tests.
• Documents the delayed batch-write race the polling addresses.

Confidence Score: 5/5

The PR appears safe to merge because the bounded polling preserves the existing spend predicates while accommodating asynchronous database convergence.

The helper repeatedly evaluates the original assertions against fresh key-info reads and still fails after a finite deadline when the expected spend state never appears.

Important Files Changed

Filename	Overview
tests/e2e/quota_management/budgets/test_budget_reset_advances_e2e.py	Adds bounded polling around two asynchronous spend assertions without introducing a concrete changed-code defect.

_{Reviews (1): Last reviewed commit: "test(e2e): poll key spend to a deadline ..." | Re-trigger Greptile}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!