From 734e633c6461dbe76af55174711f07229382581f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 18 Jul 2026 18:01:55 -0700 Subject: [PATCH 1/2] test(e2e): cover key budget_limits multi-window on personal, team, and team-member keys --- .../budgets/test_multi_window_budget_e2e.py | 71 ++++++++++++++++--- 1 file changed, 60 insertions(+), 11 deletions(-) diff --git a/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py b/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py index e1cca0c0414..9c7a65444bc 100644 --- a/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py +++ b/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py @@ -13,6 +13,10 @@ a mint-time read races the boundary; the reset job zeroes the counter in the sam pass), the key must still be refused with "over 1d budget". That check polls because enforcement's cached auth view lags the DB write; any 200 or non-budget error fails immediately. + +The blocks-then-resets check also sweeps the personal / team / team-member mint +shapes: the same budget_limits pair rides a key minted to roomy (100.0) +surroundings, so only the key's own windows can block regardless of who holds it. """ import time @@ -56,20 +60,20 @@ def _drive_to_block(client: BudgetClient, key: str) -> StreamingResponse: pytest.fail("budget never enforced before block") -@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") -def test_short_window_blocks_then_resets(client: BudgetClient, resources: ResourceManager) -> None: - key = client.generate_key( - models=[MODEL], - budget_limits=[ - BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP), - BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks - ], - ) - resources.defer(lambda: client.delete_key(key)) +def _short_roomy_limits() -> list[BudgetWindow]: + return [ + BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP), + BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks + ] + +def _assert_short_window_blocks_then_resets(client: BudgetClient, key: str) -> None: # 1. exhaust the tight window -> litellm returns budget_exceeded start = time.monotonic() - _drive_to_block(client, key) + blocked = _drive_to_block(client, key) + assert f"over {SHORT_WINDOW} budget" in blocked.body, ( + f"block must be attributed to the {SHORT_WINDOW} window, got: {blocked.body[:200]}" + ) # 2. the window resets at the next wall-clock-aligned boundary (up to a window # after start), then the reset job (~15-20s rescheduler) zeroes the spend. @@ -89,6 +93,51 @@ def test_short_window_blocks_then_resets(client: BudgetClient, resources: Resour pytest.fail(f"{WINDOW_SECONDS}s window never reset within 150s") +@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") +def test_short_window_blocks_then_resets(client: BudgetClient, resources: ResourceManager) -> None: + key = client.generate_key(models=[MODEL], budget_limits=_short_roomy_limits()) + resources.defer(lambda: client.delete_key(key)) + + _assert_short_window_blocks_then_resets(client, key) + + +def _mint_key_of_kind(client: BudgetClient, resources: ResourceManager, kind: str) -> str: + """A key carrying the tight+roomy budget_limits pair, minted to roomy (100.0) + surroundings so only the key's own windows can block.""" + match kind: + case "personal": + user_id = client.create_user(max_budget=100.0) + resources.defer(lambda: client.delete_user(user_id)) + key = client.generate_key(models=[MODEL], user_id=user_id, budget_limits=_short_roomy_limits()) + case "team": + team_id = client.create_team(alias=f"e2e-mw-team-{unique_marker()}", max_budget=100.0) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(models=[MODEL], team_id=team_id, budget_limits=_short_roomy_limits()) + case "team_member": + team_id = client.create_team(alias=f"e2e-mw-team-{unique_marker()}", max_budget=100.0) + resources.defer(lambda: client.delete_team(team_id)) + member_id = client.create_user(max_budget=100.0) + resources.defer(lambda: client.delete_user(member_id)) + client.add_team_member(team_id, member_id, max_budget_in_team=100.0) + key = client.generate_key( + models=[MODEL], team_id=team_id, user_id=member_id, budget_limits=_short_roomy_limits() + ) + case _: + pytest.fail(f"unknown key kind: {kind}") + resources.defer(lambda: client.delete_key(key)) + return key + + +@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") +@pytest.mark.parametrize("kind", ["personal", "team", "team_member"]) +def test_short_window_blocks_then_resets_across_key_kinds( + client: BudgetClient, resources: ResourceManager, kind: str +) -> None: + key = _mint_key_of_kind(client, resources, kind) + + _assert_short_window_blocks_then_resets(client, key) + + @pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") def test_long_window_blocks_after_short_window_resets(client: BudgetClient, resources: ResourceManager) -> None: key = client.generate_key( From 2ad30c04c7c007df050fe85f7f790907d96beb36 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 20 Jul 2026 10:24:24 -0700 Subject: [PATCH 2/2] refactor(e2e): group multi-window budget tests into spec classes Pure mechanical wrap: the tests move into TestKeyMultiWindowBudget and TestTeamMultiWindowBudget per the suite's lay-the-pattern-down-in-a-class convention; helpers stay module-level and bodies are unchanged besides indentation and the two signatures that would cross 120 chars --- .../budgets/test_multi_window_budget_e2e.py | 127 ++++++------- .../test_team_multi_window_budget_e2e.py | 167 +++++++++--------- 2 files changed, 150 insertions(+), 144 deletions(-) diff --git a/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py b/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py index 9c7a65444bc..80f2b6bd67d 100644 --- a/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py +++ b/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py @@ -93,14 +93,6 @@ def _assert_short_window_blocks_then_resets(client: BudgetClient, key: str) -> N pytest.fail(f"{WINDOW_SECONDS}s window never reset within 150s") -@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") -def test_short_window_blocks_then_resets(client: BudgetClient, resources: ResourceManager) -> None: - key = client.generate_key(models=[MODEL], budget_limits=_short_roomy_limits()) - resources.defer(lambda: client.delete_key(key)) - - _assert_short_window_blocks_then_resets(client, key) - - def _mint_key_of_kind(client: BudgetClient, resources: ResourceManager, kind: str) -> str: """A key carrying the tight+roomy budget_limits pair, minted to roomy (100.0) surroundings so only the key's own windows can block.""" @@ -128,68 +120,79 @@ def _mint_key_of_kind(client: BudgetClient, resources: ResourceManager, kind: st return key -@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") -@pytest.mark.parametrize("kind", ["personal", "team", "team_member"]) -def test_short_window_blocks_then_resets_across_key_kinds( - client: BudgetClient, resources: ResourceManager, kind: str -) -> None: - key = _mint_key_of_kind(client, resources, kind) +class TestKeyMultiWindowBudget: + @pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") + def test_short_window_blocks_then_resets(self, client: BudgetClient, resources: ResourceManager) -> None: + key = client.generate_key(models=[MODEL], budget_limits=_short_roomy_limits()) + resources.defer(lambda: client.delete_key(key)) - _assert_short_window_blocks_then_resets(client, key) + _assert_short_window_blocks_then_resets(client, key) -@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") -def test_long_window_blocks_after_short_window_resets(client: BudgetClient, resources: ResourceManager) -> None: - key = client.generate_key( - models=[MODEL], - budget_limits=[ - BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP), - BudgetWindow(budget_duration=LONG_WINDOW, max_budget=LONG_CAP), - ], - ) - resources.defer(lambda: client.delete_key(key)) + @pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") + @pytest.mark.parametrize("kind", ["personal", "team", "team_member"]) + def test_short_window_blocks_then_resets_across_key_kinds( + self, client: BudgetClient, resources: ResourceManager, kind: str + ) -> None: + key = _mint_key_of_kind(client, resources, kind) - # 1. drive the key to get blocked by SHORT_WINDOW, assert it's budget error - blocked = _drive_to_block(client, key) - assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}" + _assert_short_window_blocks_then_resets(client, key) - # 2. check the reset times of both budget windows after we drove to being blocked - blocked_reset_at = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW) - assert blocked_reset_at is not None, "short window missing from /key/info budget_limits" - blocked_long_reset_at = window_reset_at(client.key_budget_windows(key), LONG_WINDOW) - assert blocked_long_reset_at is not None, "long window missing from /key/info budget_limits" - # 3. poll every 5s for the SHORT_WINDOW reset time until it is past it, fails if it doesnt reset - deadline = time.monotonic() + RESET_DEADLINE_SECONDS - while time.monotonic() < deadline: - time.sleep(5) - current = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW) - if current is not None and current > blocked_reset_at: - break - else: - pytest.fail( - f"{SHORT_WINDOW} window's reset_at never advanced past {blocked_reset_at} within {RESET_DEADLINE_SECONDS}s" + @pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") + def test_long_window_blocks_after_short_window_resets( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + key = client.generate_key( + models=[MODEL], + budget_limits=[ + BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP), + BudgetWindow(budget_duration=LONG_WINDOW, max_budget=LONG_CAP), + ], ) + resources.defer(lambda: client.delete_key(key)) - # 4. short window just reset in 3, so now make a call, check that its blocked (should be blocked by LONG_WINDOW because short window reset), also make sure its budget error - deadline = time.monotonic() + RESET_DEADLINE_SECONDS - last_body = "" - while time.monotonic() < deadline: - result = _call(client, key) - if result.ok: - rolled = window_reset_at(client.key_budget_windows(key), LONG_WINDOW) != blocked_long_reset_at + # 1. drive the key to get blocked by SHORT_WINDOW, assert it's budget error + blocked = _drive_to_block(client, key) + assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}" + + # 2. check the reset times of both budget windows after we drove to being blocked + blocked_reset_at = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW) + assert blocked_reset_at is not None, "short window missing from /key/info budget_limits" + blocked_long_reset_at = window_reset_at(client.key_budget_windows(key), LONG_WINDOW) + assert blocked_long_reset_at is not None, "long window missing from /key/info budget_limits" + + # 3. poll every 5s for the SHORT_WINDOW reset time until it is past it, fails if it doesnt reset + deadline = time.monotonic() + RESET_DEADLINE_SECONDS + while time.monotonic() < deadline: + time.sleep(5) + current = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW) + if current is not None and current > blocked_reset_at: + break + else: pytest.fail( - f"{LONG_WINDOW} window failed to block after the {SHORT_WINDOW} window reset" - + (f" (the {LONG_WINDOW} window itself rolled mid-test - boundary crossed; rerun)" if rolled else "") + f"{SHORT_WINDOW} window's reset_at never advanced past {blocked_reset_at} within {RESET_DEADLINE_SECONDS}s" ) - assert is_budget_block(result), ( - f"non-budget error while waiting for {LONG_WINDOW} attribution: " - f"status={result.status_code} body={result.body[:200]}" + + # 4. short window just reset in 3, so now make a call, check that its blocked (should be blocked by LONG_WINDOW because short window reset), also make sure its budget error + deadline = time.monotonic() + RESET_DEADLINE_SECONDS + last_body = "" + while time.monotonic() < deadline: + result = _call(client, key) + if result.ok: + rolled = window_reset_at(client.key_budget_windows(key), LONG_WINDOW) != blocked_long_reset_at + pytest.fail( + f"{LONG_WINDOW} window failed to block after the {SHORT_WINDOW} window reset" + + (f" (the {LONG_WINDOW} window itself rolled mid-test - boundary crossed; rerun)" if rolled else "") + ) + assert is_budget_block(result), ( + f"non-budget error while waiting for {LONG_WINDOW} attribution: " + f"status={result.status_code} body={result.body[:200]}" + ) + if f"over {LONG_WINDOW} budget" in result.body: + return + last_body = result.body + time.sleep(5) + pytest.fail( + f"block never attributed to the {LONG_WINDOW} window within {RESET_DEADLINE_SECONDS}s: {last_body[:200]}" ) - if f"over {LONG_WINDOW} budget" in result.body: - return - last_body = result.body - time.sleep(5) - pytest.fail( - f"block never attributed to the {LONG_WINDOW} window within {RESET_DEADLINE_SECONDS}s: {last_body[:200]}" - ) diff --git a/tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py b/tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py index 1db68e6afe9..4334c036b3a 100644 --- a/tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py +++ b/tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py @@ -51,96 +51,99 @@ def _drive_to_block(client: BudgetClient, key: str) -> StreamingResponse: pytest.fail("team budget never enforced before block") -@pytest.mark.covers("quota_management.budget.team_multi_window.blocks_then_resets") -def test_team_short_window_blocks_then_resets(client: BudgetClient, resources: ResourceManager) -> None: - team_id = client.create_team( - alias=f"e2e-team-window-{unique_marker()}", - budget_limits=[ - BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=3e-6), - BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks - ], - ) - resources.defer(lambda: client.delete_team(team_id)) - key = client.generate_key(team_id=team_id, models=["claude-haiku-4-5"]) - resources.defer(lambda: client.delete_key(key)) +class TestTeamMultiWindowBudget: + @pytest.mark.covers("quota_management.budget.team_multi_window.blocks_then_resets") + def test_team_short_window_blocks_then_resets(self, client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team( + alias=f"e2e-team-window-{unique_marker()}", + budget_limits=[ + BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=3e-6), + BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks + ], + ) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(team_id=team_id, models=["claude-haiku-4-5"]) + resources.defer(lambda: client.delete_key(key)) - # 1. exhaust the tight window -> litellm returns budget_exceeded - start = time.monotonic() - _drive_to_block(client, key) + # 1. exhaust the tight window -> litellm returns budget_exceeded + start = time.monotonic() + _drive_to_block(client, key) - # 2. the window resets at the next wall-clock-aligned boundary (up to a window - # after start), then the reset job (~15-20s rescheduler) zeroes the spend. - # Allow generous headroom for that alignment + rescheduler latency; a stuck - # rescheduler is caught by the wait-loop timeout, not this elapsed bound. - deadline = time.monotonic() + 150 - while time.monotonic() < deadline: - time.sleep(5) - result = _call(client, key) - if result.ok: - elapsed = time.monotonic() - start - assert elapsed < WINDOW_SECONDS + 90, f"reset took {elapsed:.0f}s - too long for a {WINDOW_SECONDS}s window" - return - assert is_budget_block(result), f"non-budget error during reset wait: {result.body[:200]}" - pytest.fail(f"team {WINDOW_SECONDS}s window never reset within 150s") + # 2. the window resets at the next wall-clock-aligned boundary (up to a window + # after start), then the reset job (~15-20s rescheduler) zeroes the spend. + # Allow generous headroom for that alignment + rescheduler latency; a stuck + # rescheduler is caught by the wait-loop timeout, not this elapsed bound. + deadline = time.monotonic() + 150 + while time.monotonic() < deadline: + time.sleep(5) + result = _call(client, key) + if result.ok: + elapsed = time.monotonic() - start + assert elapsed < WINDOW_SECONDS + 90, f"reset took {elapsed:.0f}s - too long for a {WINDOW_SECONDS}s window" + return + assert is_budget_block(result), f"non-budget error during reset wait: {result.body[:200]}" + pytest.fail(f"team {WINDOW_SECONDS}s window never reset within 150s") -@pytest.mark.covers("quota_management.budget.team_multi_window.blocks_then_resets") -def test_team_long_window_blocks_after_short_window_resets(client: BudgetClient, resources: ResourceManager) -> None: + @pytest.mark.covers("quota_management.budget.team_multi_window.blocks_then_resets") + def test_team_long_window_blocks_after_short_window_resets( + self, client: BudgetClient, resources: ResourceManager + ) -> None: - # 0. key with a short budget window and a long budget window - team_id = client.create_team( - alias=f"e2e-team-long-window-{unique_marker()}", - budget_limits=[ - BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP), - BudgetWindow(budget_duration=LONG_WINDOW, max_budget=LONG_CAP), - ], - ) - resources.defer(lambda: client.delete_team(team_id)) - key = client.generate_key(team_id=team_id, models=["claude-haiku-4-5"]) - resources.defer(lambda: client.delete_key(key)) + # 0. key with a short budget window and a long budget window + team_id = client.create_team( + alias=f"e2e-team-long-window-{unique_marker()}", + budget_limits=[ + BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP), + BudgetWindow(budget_duration=LONG_WINDOW, max_budget=LONG_CAP), + ], + ) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(team_id=team_id, models=["claude-haiku-4-5"]) + resources.defer(lambda: client.delete_key(key)) - # 1. drive the key to being blocked, assert its blocked by budget budget_exceeded - blocked = _drive_to_block(client, key) - assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}" + # 1. drive the key to being blocked, assert its blocked by budget budget_exceeded + blocked = _drive_to_block(client, key) + assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}" - # 2. check the the teams budget windows - blocked_reset_at = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW) - assert blocked_reset_at is not None, "short window missing from /team/info budget_limits" - blocked_long_reset_at = window_reset_at(client.team_budget_windows(team_id), LONG_WINDOW) - assert blocked_long_reset_at is not None, "long window missing from /team/info budget_limits" + # 2. check the the teams budget windows + blocked_reset_at = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW) + assert blocked_reset_at is not None, "short window missing from /team/info budget_limits" + blocked_long_reset_at = window_reset_at(client.team_budget_windows(team_id), LONG_WINDOW) + assert blocked_long_reset_at is not None, "long window missing from /team/info budget_limits" - # 3. keep checking that the short budget window reset, if it doesnt within the deadline then fail - deadline = time.monotonic() + RESET_DEADLINE_SECONDS - while time.monotonic() < deadline: - time.sleep(5) - current = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW) - if current is not None and current > blocked_reset_at: - break - else: - pytest.fail( - f"team {SHORT_WINDOW} window's reset_at never advanced past " - f"{blocked_reset_at} within {RESET_DEADLINE_SECONDS}s" - ) - - # 4. short window just reset in 3, so now make another call, assert that the long window blocks the next call with budget_exceeded - deadline = time.monotonic() + RESET_DEADLINE_SECONDS - last_body = "" - while time.monotonic() < deadline: - result = _call(client, key) - if result.ok: - rolled = window_reset_at(client.team_budget_windows(team_id), LONG_WINDOW) != blocked_long_reset_at + # 3. keep checking that the short budget window reset, if it doesnt within the deadline then fail + deadline = time.monotonic() + RESET_DEADLINE_SECONDS + while time.monotonic() < deadline: + time.sleep(5) + current = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW) + if current is not None and current > blocked_reset_at: + break + else: pytest.fail( - f"team {LONG_WINDOW} window failed to block after the {SHORT_WINDOW} window reset" - + (f" (the {LONG_WINDOW} window itself rolled mid-test - boundary crossed; rerun)" if rolled else "") + f"team {SHORT_WINDOW} window's reset_at never advanced past " + f"{blocked_reset_at} within {RESET_DEADLINE_SECONDS}s" ) - assert is_budget_block(result), ( - f"non-budget error while waiting for {LONG_WINDOW} attribution: " - f"status={result.status_code} body={result.body[:200]}" + + # 4. short window just reset in 3, so now make another call, assert that the long window blocks the next call with budget_exceeded + deadline = time.monotonic() + RESET_DEADLINE_SECONDS + last_body = "" + while time.monotonic() < deadline: + result = _call(client, key) + if result.ok: + rolled = window_reset_at(client.team_budget_windows(team_id), LONG_WINDOW) != blocked_long_reset_at + pytest.fail( + f"team {LONG_WINDOW} window failed to block after the {SHORT_WINDOW} window reset" + + (f" (the {LONG_WINDOW} window itself rolled mid-test - boundary crossed; rerun)" if rolled else "") + ) + assert is_budget_block(result), ( + f"non-budget error while waiting for {LONG_WINDOW} attribution: " + f"status={result.status_code} body={result.body[:200]}" + ) + if f"over {LONG_WINDOW} budget" in result.body: + return + last_body = result.body + time.sleep(5) + pytest.fail( + f"team block never attributed to the {LONG_WINDOW} window within {RESET_DEADLINE_SECONDS}s: {last_body[:200]}" ) - if f"over {LONG_WINDOW} budget" in result.body: - return - last_body = result.body - time.sleep(5) - pytest.fail( - f"team block never attributed to the {LONG_WINDOW} window within {RESET_DEADLINE_SECONDS}s: {last_body[:200]}" - )