From e5a9f3f5d78d14511faac741a73617b767c175b9 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 17 Jul 2026 11:29:16 -0700 Subject: [PATCH] test(e2e): budget refusals are 429 for bare keys and team caps block every team key (#33632) * test(e2e): assert bare-key budget refusal is 429 and /key/info spend reaches the cap * test(e2e): keep the bare-key budget assertion to the 429 refusal shape * test(e2e): assert a team's max_budget blocks every key on the team * test(e2e): focus the team budget case on the 429 blocking behavior --- tests/e2e/CLAUDE.md | 2 +- .../coverage_registry/quota_management.yaml | 1 + .../budgets/test_budget_enforcement_e2e.py | 64 ++++++++++++++++--- 3 files changed, 58 insertions(+), 9 deletions(-) diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 0e1eafb5196..f2ca2aea437 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -131,7 +131,7 @@ Quota Management - behavior features (entity- or config-driven caps and their ac quota_management... behavior : ratelimit | budget | spend_tracking variant : rpm | tpm | priority_generous | priority_strict - key | internal_user | end_user | organization | team_member | tag + key | internal_user | end_user | organization | team | team_member | tag | model_max | soft | key_multi_window | team_multi_window | fallback | spend_counter chat_completions | stream | embeddings | cache_hit | key_rollup diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index fac266149f4..8d40a9559ea 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -7,6 +7,7 @@ - {id: quota_management.ratelimit.priority_generous.picks_under_tpm, module: quota_management, tier: P1, behavior: ratelimit, variant: priority_generous, assertions: [picks_under_tpm], exercised_on: [chat_completions, messages], source: "dynamic_rate_limiter_v3.py:36-52", rationale: "Generous mode (<80% sat) allows priority borrowing"} - {id: quota_management.ratelimit.priority_strict.picks_under_tpm, module: quota_management, tier: P1, behavior: ratelimit, variant: priority_strict, assertions: [picks_under_tpm], exercised_on: [chat_completions, messages], source: "dynamic_rate_limiter_v3.py:53-71", rationale: "Strict mode (>=80% sat) enforces priority fairness"} - {id: quota_management.budget.key.blocks_over_limit, module: quota_management, tier: P0, behavior: budget, variant: key, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A key's max_budget blocks further paid calls once spend crosses it"} +- {id: quota_management.budget.team.blocks_over_limit, module: quota_management, tier: P0, behavior: budget, variant: team, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A team's max_budget blocks every key on the team once combined spend crosses it, including keys that spent nothing themselves"} - {id: quota_management.budget.internal_user.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: internal_user, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An internal user's max_budget governs personal keys"} - {id: quota_management.budget.end_user.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: end_user, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A customer (end-user) max_budget blocks calls attributed via user="} - {id: quota_management.budget.organization.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: organization, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An organization's max_budget blocks keys under its teams"} diff --git a/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py b/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py index dbe1cfa4ea8..47cbfeb7ef0 100644 --- a/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py @@ -4,7 +4,7 @@ Each entity is an E2ECase (lifecycle.E2ECase) driven by run_case: init() creates the budgeted entity + a key, run() drives spend until a `budget_exceeded` block, teardown() deletes everything init() created (always runs, even on failure/skip). Covers the entities with no prior live coverage - internal user, end-user, -organization, team member. See BUDGET_TEST_COVERAGE_MATRIX.md. +organization, team member - plus key and team. See BUDGET_TEST_COVERAGE_MATRIX.md. A non-budget error fails hard (never a skip); if calls never get blocked, budget enforcement is broken -> fail. @@ -18,16 +18,17 @@ import pytest from budget_client import BudgetClient, is_budget_block from e2e_config import unique_marker -from e2e_http import require_successful_call +from e2e_http import StreamingResponse, require_successful_call from lifecycle import run_case pytestmark = pytest.mark.e2e -def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> None: - """Send paid calls until the entity's budget blocks one. Key/user/org/member - block within a couple calls off real-time reservation counters; the end-user - budget enforces off table spend that lands on the batch write, so it takes a - few more. A non-budget error fails hard (never a skip).""" +def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> StreamingResponse: + """Send paid calls until the entity's budget blocks one; return the blocked + response so callers can assert on its shape. Key/user/org/member block within + a couple calls off real-time reservation counters; the end-user budget + enforces off table spend that lands on the batch write, so it takes a few + more. A non-budget error fails hard (never a skip).""" for _ in range(40): result = client.chat( key, @@ -37,7 +38,7 @@ def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> user=user or None, ) if is_budget_block(result): - return + return result require_successful_call(result) time.sleep(2) pytest.fail("budget never enforced within the call budget") @@ -69,10 +70,53 @@ class _BudgetCase: class KeyBudgetCase(_BudgetCase): + """A bare key (no team_id / user_id) carrying its own max_budget, so only the + key-level budget can be the thing that blocks. The refusal must be a 429 + budget_exceeded; any other error already fails via _assert_budget_blocks.""" + def init(self) -> None: self.key = self.client.generate_key(max_budget=3e-6) self._undo.append(lambda: self.client.delete_key(self.key)) + def run(self) -> None: + blocked = _assert_budget_blocks(self.client, self.key) + assert blocked.status_code == 429, ( + f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" + ) + + +class TeamBudgetCase(_BudgetCase): + """An admin caps a whole team: two keys under a tiny-budget team, neither with + a key-level budget. Key A is driven until the team cap blocks it; key B's very + first call must then be refused too, proving the cap sits on the team, not the + key that spent. Both refusals must be 429 budget_exceeded.""" + + def init(self) -> None: + team_id = self.client.create_team( + alias=f"e2e-budget-team-{unique_marker()}", max_budget=3e-6 + ) + self._undo.append(lambda: self.client.delete_team(team_id)) + self.key = self.client.generate_key(team_id=team_id) + self._undo.append(lambda: self.client.delete_key(self.key)) + self._sibling_key = self.client.generate_key(team_id=team_id) + self._undo.append(lambda: self.client.delete_key(self._sibling_key)) + + def run(self) -> None: + blocked = _assert_budget_blocks(self.client, self.key) + assert blocked.status_code == 429, ( + f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" + ) + sibling = self.client.chat( + self._sibling_key, + "claude-haiku-4-5", + f"spend {unique_marker()}", + max_tokens=16, + ) + assert is_budget_block(sibling) and sibling.status_code == 429, ( + f"a sibling key on the capped team must get the same 429 budget_exceeded, " + f"got {sibling.status_code}: {sibling.body[:200]}" + ) + class InternalUserBudgetCase(_BudgetCase): def init(self) -> None: @@ -138,6 +182,10 @@ def _case_id(case_cls: Type[_BudgetCase]) -> str: KeyBudgetCase, marks=pytest.mark.covers("quota_management.budget.key.blocks_over_limit"), ), + pytest.param( + TeamBudgetCase, + marks=pytest.mark.covers("quota_management.budget.team.blocks_over_limit"), + ), pytest.param( InternalUserBudgetCase, marks=pytest.mark.covers("quota_management.budget.internal_user.blocks_over_limit"),