test(e2e): cover key budget_limits multi-window on personal, team, and team-member keys

[ryan-crabbe-berri]

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)

Screenshots / Proof of Fix

Originally written on top of #33832; that PR merged mid-flight, so this branch is rebased onto staging past the squash merge

All runs are against a live proxy on localhost:4000 running staging code with a throwaway Postgres and Redis, the fast budget rescheduler, and real provider calls (gpt-5.5 and claude-haiku-4-5 are backed by groq/llama-3.3-70b-versatile locally since that is the only live provider key on this machine)

The new sweep ran five consecutive times at commit 734e633 with cooldowns: 3 passed in 221.61s, 239.20s, 194.12s, 249.34s, 259.16s; zero failures (the spread is 30s wall-clock window alignment)

$ python -m pytest "tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py::TestKeyMultiWindowBudget::test_short_window_blocks_then_resets_across_key_kinds" -v
... test_short_window_blocks_then_resets_across_key_kinds[personal] PASSED [ 33%]
... test_short_window_blocks_then_resets_across_key_kinds[team] PASSED [ 66%]
... test_short_window_blocks_then_resets_across_key_kinds[team_member] PASSED [100%]

Proof restamped after the class-wrap commit at head 2ad30c0: three consecutive full-file runs of each touched file, zero failures

$ for run in 1 2 3; do python -m pytest tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py -q | tail -1; sleep 15; done
5 passed in 412.21s (0:06:52)
5 passed in 405.72s (0:06:45)
5 passed in 399.98s (0:06:39)
$ for run in 1 2 3; do python -m pytest tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py -q | tail -1; sleep 15; done
2 passed in 139.06s (0:02:19)
2 passed in 169.30s (0:02:49)
2 passed in 128.90s (0:02:08)

basedpyright tests/e2e reports 0 errors and python -m coverage_registry.collector --strict passes

Type

✅ Test

Changes

Two commits. The first covers the "Key budget_limits (multi-window) blocks + resets" row of the budgets coverage matrix for the three unexercised key kinds. The existing test_short_window_blocks_then_resets proves a bare key's tight 30s window trips and self-heals while a roomy 1m window rides along; the new parametrized test_short_window_blocks_then_resets_across_key_kinds proves the same key-level windows still enforce and reset when the key is minted to a person (user_id), a team (team_id), or a team membership (both), each with roomy (100.0) surrounding budgets so only the key's own windows can block. The shared body is hoisted into _assert_short_window_blocks_then_resets and now also asserts the block is attributed to the 30s window ("over 30s budget" in the refusal body, mirroring the long-window test's attribution check), so a future change that lets a surrounding budget produce the block fails instead of passing silently

The second commit is a pure mechanical wrap: the tests of this file and its team sibling move into TestKeyMultiWindowBudget and TestTeamMultiWindowBudget spec classes, per the suite's lay-the-pattern-down-in-a-class convention (these two files were the only budgets files with loose module-level tests). Helpers stay module-level; bodies are unchanged besides indentation, self, and the two def signatures that would cross 120 chars. Node ids gain the class prefix

All three parametrized cases intentionally mark the existing registry cell quota_management.budget.key_multi_window.blocks_then_resets (shared with the bare-key test): the matrix tracks the per-key-kind split while the registry keeps one row per budget level x behavior. The long-window direction stays covered on the bare-key shape by #33832's test and is not repeated per key kind

QA runbook

• TestKeyMultiWindowBudget::test_short_window_blocks_then_resets_across_key_kinds[personal] creates a user with max_budget 100, mints that user a personal key whose budget_limits are 30s at 1e-9 and 1m at 1.0, drives /chat/completions until refused, asserts the refusal names the 30s window, then polls until calls flow again within the elapsed bound, holding every refusal during the wait to a budget block; teardown deletes the key, then the user
• TestKeyMultiWindowBudget::test_short_window_blocks_then_resets_across_key_kinds[team] does the same with a team at max_budget 100 and a team key carrying the window pair; teardown deletes the key, then the team
• TestKeyMultiWindowBudget::test_short_window_blocks_then_resets_across_key_kinds[team_member] does the same with a team and user at max_budget 100 each, the membership at max_budget_in_team 100, and a team-member key carrying the window pair; teardown deletes the key, the user, then the team
• TestKeyMultiWindowBudget::test_short_window_blocks_then_resets still passes through the hoisted helper with its original bare-key mint shape
• TestKeyMultiWindowBudget::test_long_window_blocks_after_short_window_resets and both TestTeamMultiWindowBudget tests still pass unchanged inside their new classes
• python -m coverage_registry.collector --strict passes with the parametrized cases sharing quota_management.budget.key_multi_window.blocks_then_resets

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 PR extends e2e coverage for key-level budget_limits (multi-window) by adding a parametrized test that exercises the same tight-30s-window enforcement and reset cycle across three key mint shapes: personal (user_id), team (team_id), and team-member (both). It also wraps the existing module-level tests in both budget files into spec classes per the suite's convention.

• Extracts _short_roomy_limits(), _assert_short_window_blocks_then_resets(), and _mint_key_of_kind() helpers; the shared assert now additionally checks that the refusal body names the 30s window (a net strengthening over the original test).
• test_team_multi_window_budget_e2e.py is a mechanical class-wrap only — no assertion or logic changes.

Confidence Score: 5/5

Safe to merge — test-only changes with no production code touched

Both files are e2e tests living entirely under tests/e2e/. The refactoring is mechanical (class wrap, helper extraction) and does not remove or weaken any existing assertions — it adds one (block attribution must name the 30s window). The new parametrized test follows the same flow as the pre-existing bare-key test but mints three different key contexts. Cleanup ordering via resources.defer is consistent with the LIFO pattern already used throughout the suite. No production code is changed.

No files require special attention

Important Files Changed

Filename	Overview
tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py	Adds parametrized test covering personal/team/team-member key kinds, extracts shared helpers, wraps tests into class, and strengthens attribution assertion; no logic regressions found
tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py	Pure mechanical wrap of existing tests into TestTeamMultiWindowBudget class; no logic changes

_{Reviews (2): Last reviewed commit: "refactor(e2e): group multi-window budget..." | Re-trigger Greptile}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_e2e_key_multi_window_key_kinds (2ad30c0) with litellm_internal_staging (3fcd19d)}

[ryan-crabbe-berri]

@greptileai re review