From eb8870065bd96061d69445f2ec5b3f1ea6fd891c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 30 Jul 2026 14:09:50 -0700 Subject: [PATCH] test(e2e): align budget e2e with the team-key budget hierarchy (#35276) #35271 restored the hierarchy where a team-scoped key is governed by the team and team-member budgets only; the owner's personal max_budget applies to their personal keys. Three places in the e2e suite still encoded the old direction and would fail against a proxy built from staging. test_user_budget_enforced_across_all_their_keys asserted that the owner's team-member key is refused once their personal budget is exhausted. It now asserts only the personal keys are refused, and keeps the team key as the control that must keep serving, which pins the restored direction instead of leaving it unasserted. Renamed to match what it now covers. test_team_member_key_user_budget_resets_after_window drove a team key to a block off the owner's personal budget, so nothing can block it any more and _drive_to_block could never succeed. Its premise is gone rather than moved, so it is removed; the sibling personal-key test still covers quota_management.budget.internal_user.resets_after_window. The registry rationale for that row dropped its "and team-member keys" clause for the same reason. --- .../coverage_registry/quota_management.yaml | 2 +- .../budgets/test_budget_enforcement_e2e.py | 24 +++++++++++++------ .../budgets/test_budget_reset_e2e.py | 15 ------------ 3 files changed, 18 insertions(+), 23 deletions(-) diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index a8d0749cd8d..eb620395c46 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -23,7 +23,7 @@ - {id: quota_management.budget.key.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: key, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_duration zeroes key spend after the window; a blocked key serves again"} - {id: quota_management.budget.team.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: team, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_duration zeroes a team's spend after the window; every key on the team serves again"} - {id: quota_management.budget.organization.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: organization, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "An org budget resets after its window; keys under the org serve again"} -- {id: quota_management.budget.internal_user.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: internal_user, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "An internal user's budget resets after its window; their personal and team-member keys serve again"} +- {id: quota_management.budget.internal_user.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: internal_user, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "An internal user's budget resets after its window; their personal keys serve again"} - {id: quota_management.budget.team_member.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "Member per-team budget reset keeps advancing window after window"} - {id: quota_management.budget.key_multi_window.blocks_then_resets, module: quota_management, tier: P1, behavior: budget, variant: key_multi_window, assertions: [blocks_then_resets], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_limits enforce within a short window and serve again in the next"} - {id: quota_management.budget.key_multi_window.resets_windows_independently, module: quota_management, tier: P2, behavior: budget, variant: key_multi_window, assertions: [resets_windows_independently], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "Each window of a multi-window budget resets on its own schedule"} 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 8b93afb4752..918739863ce 100644 --- a/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py @@ -79,9 +79,14 @@ class TestBudgetBlocksPerLevel: ) @pytest.mark.covers("quota_management.budget.internal_user.blocks_over_limit") - def test_user_budget_enforced_across_all_their_keys( + def test_user_budget_enforced_across_their_personal_keys( self, client: BudgetClient, resources: ResourceManager ) -> None: + """A user's max_budget follows the person across their personal keys, so a + second untouched key is not a fresh allowance. It stops at the team + boundary: the same user's team-scoped key is governed by the team and + team-member budgets, both uncapped here, so it is the control that must + keep serving while the personal keys are refused.""" user_id = client.create_user(max_budget=TINY_CAP) resources.defer(lambda: client.delete_user(user_id)) first_key = client.generate_key(user_id=user_id) @@ -95,12 +100,17 @@ class TestBudgetBlocksPerLevel: resources.defer(lambda: client.delete_key(team_key)) _assert_blocked_429(client, first_key) - for label, key in (("second personal key", second_key), ("team-member key", team_key)): - result = _chat(client, key) - assert is_budget_block(result) and result.status_code == 429, ( - f"the {label} of a user over budget must get the same 429 budget_exceeded, " - f"got {result.status_code}: {result.body[:200]}" - ) + second = _chat(client, second_key) + assert is_budget_block(second) and second.status_code == 429, ( + f"the second personal key of a user over budget must get the same 429 budget_exceeded, " + f"got {second.status_code}: {second.body[:200]}" + ) + team_result = _chat(client, team_key) + assert not is_budget_block(team_result), ( + f"the team-scoped key of a user over their personal budget must keep serving; " + f"got {team_result.status_code}: {team_result.body[:200]}" + ) + require_successful_call(team_result) @pytest.mark.covers("quota_management.budget.end_user.blocks_over_limit") def test_end_user_budget_blocks_attributed_calls( diff --git a/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py b/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py index 793b22a47c7..b7b7f269c47 100644 --- a/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py @@ -102,21 +102,6 @@ class TestBudgetResetPerLevel: _drive_to_block(client, key) _poll_until_serves_again(client, key) - @pytest.mark.covers("quota_management.budget.internal_user.resets_after_window") - def test_team_member_key_user_budget_resets_after_window( - self, client: BudgetClient, resources: ResourceManager - ) -> None: - user_id = client.create_user(max_budget=TINY_CAP, budget_duration=WINDOW) - resources.defer(lambda: client.delete_user(user_id)) - team_id = client.create_team(alias=f"e2e-user-team-reset-{unique_marker()}") - resources.defer(lambda: client.delete_team(team_id)) - client.add_team_member(team_id, user_id, max_budget_in_team=100.0) - key = client.generate_key(team_id=team_id, user_id=user_id) - resources.defer(lambda: client.delete_key(key)) - - _drive_to_block(client, key) - _poll_until_serves_again(client, key) - class TestKeyBudgetResetAcrossKeyKinds: """The tiny max_budget and its 30s window sit on the key itself while the user,