From 9f6b3d24f72cc218cd88671552cfd330fd5e2765 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 22 Jul 2026 17:48:18 -0700 Subject: [PATCH] test(e2e): cover key max_budget blocks on personal, team, and team-member keys (#33895) * test(e2e): cover key max_budget blocks on personal, team, and team-member keys * refactor(e2e): convert budget enforcement cases to the resources-fixture pattern The E2ECase class pattern existed only in this file; every other suite uses plain pytest tests with the resources fixture. Rewrites the nine cases as two spec classes and removes the now-dead E2ECase protocol and run_case driver from lifecycle.py --- tests/e2e/conftest.py | 6 +- tests/e2e/lifecycle.py | 42 +-- .../budgets/test_budget_enforcement_e2e.py | 349 +++++++----------- 3 files changed, 147 insertions(+), 250 deletions(-) diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 7b5ad3551b8..eff3b4ddf58 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -5,9 +5,9 @@ answers or when credentials/env are missing; they never skip. Pure unit coverage of the harness itself carries no `e2e` marker and runs regardless of whether a proxy is up. -Lifecycle: the `resources` fixture maps the init -> run -> teardown contract -(lifecycle.E2ECase) onto pytest - setup is init(), the test body is run(), and -teardown deletes every resource the test created on the long-lived proxy. +Lifecycle: the `resources` fixture hands each test a lifecycle.ResourceManager - +the test registers a cleanup for every resource it creates, and the fixture's +teardown deletes them all on the long-lived proxy, even when the test fails. Each suite provides its own `client` fixture (a lifecycle.ResourceClient); these shared fixtures build on it. diff --git a/tests/e2e/lifecycle.py b/tests/e2e/lifecycle.py index c0de074f0d2..4ef25509905 100644 --- a/tests/e2e/lifecycle.py +++ b/tests/e2e/lifecycle.py @@ -1,13 +1,11 @@ -"""Lifecycle contract and resource cleanup for stateful e2e tests. +"""Resource cleanup for stateful e2e tests. Shared by every e2e suite under tests/e2e/. The proxy under test is long-lived and never reset between tests, so anything a test creates (keys, customers, teams, orgs, users, guardrails, budgets, ...) persists unless -explicitly deleted. Every check follows an init -> run -> teardown lifecycle; -teardown releases each resource init() created, even when run() raises. - -In pytest terms (see conftest.py): the `resources` fixture's setup is init(), -the test body is run(), and the fixture's teardown is teardown(). +explicitly deleted. The `resources` fixture (see conftest.py) hands each test a +ResourceManager; the test registers a cleanup for every resource it creates, and +the fixture's teardown releases them all even when the test body raises. """ from dataclasses import dataclass, field @@ -17,38 +15,6 @@ from proxy_client import ProxyClient from models import KeyGenerateBody -@runtime_checkable -class E2ECase(Protocol): - """A stateful e2e check run against a long-lived proxy. - - init() acquires resources, run() exercises behaviour and asserts, teardown() - releases everything init() created. teardown() must run even if init() fails - partway or run() raises. - """ - - def init(self) -> None: ... - - def run(self) -> None: ... - - def teardown(self) -> None: ... - - -def run_case(case: E2ECase) -> None: - """Drive a case through its lifecycle: init -> run -> teardown. - - teardown always runs - even when init() fails partway or run() raises (or - skips) - so resources the case already registered on the long-lived proxy are - released. init() is inside the try because cases register cleanups - progressively (e.g. create team, then user, then key), and a failure after - the first creation must still release what came before. - """ - try: - case.init() - case.run() - finally: - case.teardown() - - @runtime_checkable class ResourceClient(Protocol): """Proxy operations the convenience creators use. Resource types without a 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 c4ad0c38f31..8b93afb4752 100644 --- a/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py @@ -1,28 +1,36 @@ """Live e2e: a tiny max_budget on an entity actually blocks requests. -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 - plus key and team. See BUDGET_TEST_COVERAGE_MATRIX.md. +One test per budget level (key, team, internal user, end-user, organization, +team member): put the tiny cap on that level, drive spend until a +`budget_exceeded` block, and where a cap could be confused with a neighbor, +prove isolation with an uncapped control key that must keep serving. The +capped-key sweep proves the key's own max_budget blocks across mint shapes +(personal, team, team-member) with roomy surroundings, so the key-level cap is +provably the blocker no matter who the key was minted to. A non-budget error fails hard (never a skip); if calls never get blocked, budget enforcement is broken -> fail. """ import time -from dataclasses import dataclass, field -from typing import Callable, List, Type import pytest from budget_client import BudgetClient, is_budget_block from e2e_config import unique_marker from e2e_http import StreamingResponse, require_successful_call -from lifecycle import run_case +from lifecycle import ResourceManager pytestmark = pytest.mark.e2e +TINY_CAP = 3e-6 +ROOMY_CAP = 100.0 + + +def _chat(client: BudgetClient, key: str, *, user: str | None = None) -> StreamingResponse: + return client.chat(key, "claude-haiku-4-5", f"spend {unique_marker()}", max_tokens=16, user=user) + + 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 @@ -30,13 +38,7 @@ def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> 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, - "claude-haiku-4-5", - f"spend {unique_marker()}", - max_tokens=16, - user=user or None, - ) + result = _chat(client, key, user=user or None) if is_budget_block(result): return result require_successful_call(result) @@ -44,225 +46,154 @@ def _assert_budget_blocks(client: BudgetClient, key: str, *, user: str = "") -> pytest.fail("budget never enforced within the call budget") -@dataclass -class _BudgetCase: - """Base E2ECase: a key under some budgeted entity must get blocked. - - Subclasses set up the budgeted entity in init() and register every created id - in `_undo` (run LIFO in teardown so a key is deleted before its team/org). - """ - - client: BudgetClient - key: str = "" - _undo: List[Callable[[], None]] = field( - default_factory=list - ) # mutable-ok: per-case teardown registry - - def init(self) -> None: - raise NotImplementedError - - def run(self) -> None: - _assert_budget_blocks(self.client, self.key) - - def teardown(self) -> None: - for undo in reversed(self._undo): - undo() +def _assert_blocked_429(client: BudgetClient, key: str) -> StreamingResponse: + blocked = _assert_budget_blocks(client, key) + assert blocked.status_code == 429, ( + f"budget refusal must be 429, got {blocked.status_code}: {blocked.body[:200]}" + ) + return blocked -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.""" +class TestBudgetBlocksPerLevel: + @pytest.mark.covers("quota_management.budget.key.blocks_over_limit") + def test_bare_key_blocks_over_its_own_budget(self, client: BudgetClient, resources: ResourceManager) -> None: + key = client.generate_key(max_budget=TINY_CAP) + resources.defer(lambda: client.delete_key(key)) - def init(self) -> None: - self.key = self.client.generate_key(max_budget=3e-6) - self._undo.append(lambda: self.client.delete_key(self.key)) + _assert_blocked_429(client, 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]}" - ) + @pytest.mark.covers("quota_management.budget.team.blocks_over_limit") + def test_team_budget_blocks_every_team_key(self, client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team(alias=f"e2e-budget-team-{unique_marker()}", max_budget=TINY_CAP) + resources.defer(lambda: client.delete_team(team_id)) + spender_key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(spender_key)) + sibling_key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(sibling_key)) - -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_blocked_429(client, spender_key) + sibling = _chat(client, sibling_key) 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]}" ) + @pytest.mark.covers("quota_management.budget.internal_user.blocks_over_limit") + def test_user_budget_enforced_across_all_their_keys( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + 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) + resources.defer(lambda: client.delete_key(first_key)) + second_key = client.generate_key(user_id=user_id) + resources.defer(lambda: client.delete_key(second_key)) + team_id = client.create_team(alias=f"e2e-budget-team-{unique_marker()}") + resources.defer(lambda: client.delete_team(team_id)) + client.add_team_member(team_id, user_id) + team_key = client.generate_key(team_id=team_id, user_id=user_id) + resources.defer(lambda: client.delete_key(team_key)) -class InternalUserBudgetCase(_BudgetCase): - """A user's max_budget follows the person, not the key. The capped user holds - two personal keys (no team, no key budgets) plus a team-member key on an - uncapped team; once the first personal key is refused, the other two must be - refused as well - a second key is not a fresh allowance, and since #32005 the - user budget draws down team keys too. All refusals must be 429 budget_exceeded.""" - - def init(self) -> None: - user_id = self.client.create_user(max_budget=3e-6) - self._undo.append(lambda: self.client.delete_user(user_id)) - self.key = self.client.generate_key(user_id=user_id) - self._undo.append(lambda: self.client.delete_key(self.key)) - self._second_key = self.client.generate_key(user_id=user_id) - self._undo.append(lambda: self.client.delete_key(self._second_key)) - team_id = self.client.create_team(alias=f"e2e-budget-team-{unique_marker()}") - self._undo.append(lambda: self.client.delete_team(team_id)) - self.client.add_team_member(team_id, user_id) - self._team_key = self.client.generate_key(team_id=team_id, user_id=user_id) - self._undo.append(lambda: self.client.delete_key(self._team_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]}" - ) - for label, key in (("second personal key", self._second_key), ("team-member key", self._team_key)): - result = self.client.chat(key, "claude-haiku-4-5", f"spend {unique_marker()}", max_tokens=16) + _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]}" ) - -class EndUserBudgetCase(_BudgetCase): - def init(self) -> None: + @pytest.mark.covers("quota_management.budget.end_user.blocks_over_limit") + def test_end_user_budget_blocks_attributed_calls( + self, client: BudgetClient, resources: ResourceManager + ) -> None: customer = f"e2e-budget-cust-{unique_marker()}" - self.client.create_customer(customer, max_budget=3e-6) - self._undo.append(lambda: self.client.delete_customers([customer])) - self.key = self.client.generate_key(models=["claude-haiku-4-5"]) - self._undo.append(lambda: self.client.delete_key(self.key)) - self._customer = customer + client.create_customer(customer, max_budget=TINY_CAP) + resources.defer(lambda: client.delete_customers([customer])) + key = client.generate_key(models=["claude-haiku-4-5"]) + resources.defer(lambda: client.delete_key(key)) - def run(self) -> None: - _assert_budget_blocks(self.client, self.key, user=self._customer) + _assert_budget_blocks(client, key, user=customer) + @pytest.mark.covers("quota_management.budget.organization.blocks_over_limit") + def test_org_budget_blocks_keys_under_it(self, client: BudgetClient, resources: ResourceManager) -> None: + org_id = client.create_org(max_budget=TINY_CAP, alias=f"e2e-budget-org-{unique_marker()}") + resources.defer(lambda: client.delete_org(org_id)) + team_id = client.create_team(alias=f"e2e-budget-team-{unique_marker()}", organization_id=org_id) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(key)) -class OrganizationBudgetCase(_BudgetCase): - """Org carries the tiny budget; the team under it and the key carry none, so - the org is the only entity that can block (the historically weak link). The - refusal must be a 429 budget_exceeded that names the org as the blocker.""" - - def init(self) -> None: - self._org_id = self.client.create_org( - max_budget=3e-6, alias=f"e2e-budget-org-{unique_marker()}" - ) - self._undo.append(lambda: self.client.delete_org(self._org_id)) - team_id = self.client.create_team( - alias=f"e2e-budget-team-{unique_marker()}", organization_id=self._org_id - ) - 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)) - - 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]}" - ) - assert f"Organization={self._org_id}" in blocked.body, ( + blocked = _assert_blocked_429(client, key) + assert f"Organization={org_id}" in blocked.body, ( f"refusal must name the org as the blocker, got: {blocked.body[:200]}" ) + @pytest.mark.covers("quota_management.budget.team_member.blocks_over_limit") + def test_member_budget_blocks_without_touching_teammates( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + team_id = client.create_team(alias=f"e2e-budget-team-{unique_marker()}", max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_team(team_id)) + member_id = client.create_user(max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_user(member_id)) + client.add_team_member(team_id, member_id, max_budget_in_team=TINY_CAP) + member_key = client.generate_key(team_id=team_id, user_id=member_id) + resources.defer(lambda: client.delete_key(member_key)) + teammate_id = client.create_user(max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_user(teammate_id)) + client.add_team_member(team_id, teammate_id) + teammate_key = client.generate_key(team_id=team_id, user_id=teammate_id) + resources.defer(lambda: client.delete_key(teammate_key)) -class TeamMemberBudgetCase(_BudgetCase): - """Member A's per-team budget is tiny while the team and both members' user - budgets are roomy (100.0), so the only cap that can trip is A's: a block - proves member-level enforcement and must be a 429 budget_exceeded. Teammate - B, uncapped on the same team, must keep serving after A is cut off, proving - the member cap does not leak onto the team or its members.""" - - def init(self) -> None: - self._team_id = self.client.create_team( - alias=f"e2e-budget-team-{unique_marker()}", max_budget=100.0 - ) - self._undo.append(lambda: self.client.delete_team(self._team_id)) - self._member_id = self.client.create_user(max_budget=100.0) - self._undo.append(lambda: self.client.delete_user(self._member_id)) - self.client.add_team_member(self._team_id, self._member_id, max_budget_in_team=3e-6) - self.key = self.client.generate_key(team_id=self._team_id, user_id=self._member_id) - self._undo.append(lambda: self.client.delete_key(self.key)) - teammate_id = self.client.create_user(max_budget=100.0) - self._undo.append(lambda: self.client.delete_user(teammate_id)) - self.client.add_team_member(self._team_id, teammate_id) - self._teammate_key = self.client.generate_key(team_id=self._team_id, user_id=teammate_id) - self._undo.append(lambda: self.client.delete_key(self._teammate_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]}" - ) - teammate = self.client.chat( - self._teammate_key, - "claude-haiku-4-5", - f"spend {unique_marker()}", - max_tokens=16, - ) - require_successful_call(teammate) + _assert_blocked_429(client, member_key) + require_successful_call(_chat(client, teammate_key)) -def _case_id(case_cls: Type[_BudgetCase]) -> str: - return case_cls.__name__ +class TestKeyBudgetBlocksAcrossKeyKinds: + """The tiny max_budget sits on the key itself while every budget around it + (user / team / membership) is roomy, so only the key-level cap can block; the + uncapped control key minted to the same surroundings must keep serving after + the capped key is refused, proving nothing around the key was the blocker.""" + @pytest.mark.covers("quota_management.budget.key.blocks_over_limit") + def test_personal_key_blocks_over_its_own_budget( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + user_id = client.create_user(max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_user(user_id)) + capped_key = client.generate_key(user_id=user_id, max_budget=TINY_CAP) + resources.defer(lambda: client.delete_key(capped_key)) + control_key = client.generate_key(user_id=user_id) + resources.defer(lambda: client.delete_key(control_key)) -@pytest.mark.parametrize( - "case_cls", - [ - pytest.param( - 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"), - ), - pytest.param( - EndUserBudgetCase, - marks=pytest.mark.covers("quota_management.budget.end_user.blocks_over_limit"), - ), - pytest.param( - OrganizationBudgetCase, - marks=pytest.mark.covers("quota_management.budget.organization.blocks_over_limit"), - ), - pytest.param( - TeamMemberBudgetCase, - marks=pytest.mark.covers("quota_management.budget.team_member.blocks_over_limit"), - ), - ], - ids=_case_id, -) -def test_budget_enforcement( - client: BudgetClient, case_cls: Type[_BudgetCase] -) -> None: - run_case(case_cls(client)) + _assert_blocked_429(client, capped_key) + require_successful_call(_chat(client, control_key)) + + @pytest.mark.covers("quota_management.budget.key.blocks_over_limit") + def test_team_key_blocks_over_its_own_budget(self, client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team(alias=f"e2e-key-cap-team-{unique_marker()}", max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_team(team_id)) + capped_key = client.generate_key(team_id=team_id, max_budget=TINY_CAP) + resources.defer(lambda: client.delete_key(capped_key)) + control_key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(control_key)) + + _assert_blocked_429(client, capped_key) + require_successful_call(_chat(client, control_key)) + + @pytest.mark.covers("quota_management.budget.key.blocks_over_limit") + def test_team_member_key_blocks_over_its_own_budget( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + team_id = client.create_team(alias=f"e2e-key-cap-team-{unique_marker()}", max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_team(team_id)) + member_id = client.create_user(max_budget=ROOMY_CAP) + resources.defer(lambda: client.delete_user(member_id)) + client.add_team_member(team_id, member_id, max_budget_in_team=ROOMY_CAP) + capped_key = client.generate_key(team_id=team_id, user_id=member_id, max_budget=TINY_CAP) + resources.defer(lambda: client.delete_key(capped_key)) + control_key = client.generate_key(team_id=team_id, user_id=member_id) + resources.defer(lambda: client.delete_key(control_key)) + + _assert_blocked_429(client, capped_key) + require_successful_call(_chat(client, control_key))