diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 0d61f48703d..633351ca97c 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -18,6 +18,9 @@ - {id: quota_management.budget.model_max.isolates_per_model, module: quota_management, tier: P1, behavior: budget, variant: model_max, assertions: [isolates_per_model], exercised_on: [chat_completions], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "model_max_budget caps one model without touching a sibling's budget"} - {id: quota_management.budget.soft.alerts_without_blocking, module: quota_management, tier: P1, behavior: budget, variant: soft, assertions: [alerts_without_blocking], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "soft_budget alerts but never blocks traffic"} - {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.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/budget_client.py b/tests/e2e/quota_management/budgets/budget_client.py index 01e4d63c1c3..c2f5dcdd57c 100644 --- a/tests/e2e/quota_management/budgets/budget_client.py +++ b/tests/e2e/quota_management/budgets/budget_client.py @@ -33,6 +33,7 @@ _TEAM_READY_SLEEP_SECONDS = 0.4 class UserNewBody(BaseModel): max_budget: float + budget_duration: str | None = None class UserNewResponse(BaseModel): @@ -64,6 +65,7 @@ class CustomerNewBody(BaseModel): class OrgNewBody(BaseModel): organization_alias: str max_budget: float + budget_duration: str | None = None class OrgNewResponse(BaseModel): @@ -74,6 +76,14 @@ class OrgDeleteBody(BaseModel): organization_ids: list[str] +class OrgInfoParams(BaseModel): + organization_id: str + + +class OrgInfoResponse(BaseModel): + budget_id: str | None = None + + class TeamMember(BaseModel): role: str user_id: str @@ -82,6 +92,7 @@ class TeamMember(BaseModel): class TeamNewBody(BaseModel): team_alias: str max_budget: float | None = None + budget_duration: str | None = None organization_id: str | None = None budget_limits: list[BudgetWindow] | None = None @@ -257,12 +268,12 @@ class BudgetClient: # ---- internal user -------------------------------------------------- - def create_user(self, *, max_budget: float) -> str: + def create_user(self, *, max_budget: float, budget_duration: str | None = None) -> str: return unwrap( self.gateway.transport.post( "/user/new", headers=self.gateway.transport.master, - json=UserNewBody(max_budget=max_budget), + json=UserNewBody(max_budget=max_budget, budget_duration=budget_duration), response_type=UserNewResponse, ) ).user_id @@ -301,16 +312,36 @@ class BudgetClient: # ---- organization --------------------------------------------------- - def create_org(self, *, max_budget: float, alias: str) -> str: + def create_org(self, *, max_budget: float, alias: str, budget_duration: str | None = None) -> str: return unwrap( self.gateway.transport.post( "/organization/new", headers=self.gateway.transport.master, - json=OrgNewBody(organization_alias=alias, max_budget=max_budget), + json=OrgNewBody( + organization_alias=alias, + max_budget=max_budget, + budget_duration=budget_duration, + ), response_type=OrgNewResponse, ) ).organization_id + def org_budget_id(self, org_id: str) -> str | None: + """The id of the budget row backing an org; its budget_reset_at is read via + budget_info (LIT-4570: /organization/new stores budget_duration without + scheduling budget_reset_at, so the reset job's first tick schedules it).""" + result = self.gateway.transport.get( + "/organization/info", + headers=self.gateway.transport.master, + params=OrgInfoParams(organization_id=org_id), + response_type=OrgInfoResponse, + ) + match result: + case Success(data=data): + return data.budget_id + case _: + return None + def delete_org(self, org_id: str) -> None: _ = self.gateway.transport.delete( "/organization/delete", @@ -326,6 +357,7 @@ class BudgetClient: *, alias: str, max_budget: float | None = None, + budget_duration: str | None = None, organization_id: str | None = None, budget_limits: list[BudgetWindow] | None = None, ) -> str: @@ -336,6 +368,7 @@ class BudgetClient: json=TeamNewBody( team_alias=alias, max_budget=max_budget, + budget_duration=budget_duration, organization_id=organization_id, budget_limits=budget_limits, ), 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 bdbee027f28..b57ad23ebf5 100644 --- a/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_reset_e2e.py @@ -1,12 +1,4 @@ -"""Live e2e: a key budget resets (zeroes spend) after its budget_duration. - -Short budget_duration (30s) + the fast-rescheduled reset job: a key blocked for -exceeding its max_budget starts succeeding again once the duration elapses and the -reset job zeroes key.spend. Closes the reset-zeroing gap in -BUDGET_TEST_COVERAGE_MATRIX.md (reset_budget_for_litellm_keys), which the unit -suite covers but no live test did - distinct from the per-window reset in -test_multi_window_budget_e2e.py. -""" +"""Live e2e: an entity blocked over its max_budget serves again after its budget_duration window.""" import time @@ -19,42 +11,107 @@ from lifecycle import ResourceManager pytestmark = pytest.mark.e2e +TINY_CAP = 3e-6 +WINDOW = "30s" +RESET_DEADLINE_SECONDS = 150 + def _call(client: BudgetClient, key: str): - return client.chat( - key, "claude-haiku-4-5", f"reset {unique_marker()}", max_tokens=16 - ) + return client.chat(key, "claude-haiku-4-5", f"reset {unique_marker()}", max_tokens=16) -@pytest.mark.covers("quota_management.budget.key.resets_after_window") -def test_key_budget_resets_after_duration( - client: BudgetClient, resources: ResourceManager -) -> None: - key = client.generate_key(max_budget=3e-6, budget_duration="30s") - resources.defer(lambda: client.delete_key(key)) - - # 1. exceed the budget -> litellm returns budget_exceeded - blocked = False - for _ in range(20): +def _drive_to_block(client: BudgetClient, key: str) -> None: + """Spend until the cap blocks a call, staying under one window so the block + is observed before the reset job can fire; fail hard if enforcement never trips.""" + for _ in range(12): result = _call(client, key) if is_budget_block(result): - blocked = True - break + return require_successful_call(result) time.sleep(2) - assert blocked, "key budget never enforced" + pytest.fail("budget never enforced before the window could reset") - # 2. once the 30s duration elapses + the reset job runs, key.spend zeroes and - # calls flow again. The window is wall-clock-aligned, so the reset lands up to - # a window later, then the rescheduler (~15-20s) zeroes the spend; allow - # generous headroom over that. A stuck rescheduler is caught by the wait-loop - # timeout, not this elapsed bound. - start = time.monotonic() - while time.monotonic() < start + 150: + +def _poll_until_serves_again(client: BudgetClient, key: str) -> None: + """Poll past the window until the blocked key serves again; every refusal must + stay a budget block, so a crashed reset path or provider error fails loudly.""" + deadline = time.monotonic() + RESET_DEADLINE_SECONDS + while time.monotonic() < deadline: time.sleep(5) result = _call(client, key) if result.ok: - assert time.monotonic() - start < 120, "reset too slow for a 30s budget" return - assert is_budget_block(result), f"non-budget error: {result.body[:200]}" - pytest.fail("key budget never reset within 150s") + if not is_budget_block(result): + pytest.fail(f"non-budget error during reset wait: HTTP {result.status_code}: {result.body[:200]}") + pytest.fail(f"budget never reset within {RESET_DEADLINE_SECONDS}s") + + +class TestBudgetResetDiagonal: + @pytest.mark.covers("quota_management.budget.key.resets_after_window") + def test_bare_key_budget_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None: + key = client.generate_key(max_budget=TINY_CAP, budget_duration=WINDOW) + resources.defer(lambda: client.delete_key(key)) + + _drive_to_block(client, key) + _poll_until_serves_again(client, key) + + @pytest.mark.covers("quota_management.budget.team.resets_after_window") + def test_team_budget_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team( + alias=f"e2e-team-reset-{unique_marker()}", max_budget=TINY_CAP, budget_duration=WINDOW + ) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(team_id=team_id) + resources.defer(lambda: client.delete_key(key)) + + _drive_to_block(client, key) + _poll_until_serves_again(client, key) + + @pytest.mark.covers("quota_management.budget.organization.resets_after_window") + def test_org_budget_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None: + org_id = client.create_org( + max_budget=TINY_CAP, alias=f"e2e-org-reset-{unique_marker()}", budget_duration=WINDOW + ) + resources.defer(lambda: client.delete_org(org_id)) + team_id = client.create_team(alias=f"e2e-org-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)) + + budget_id = client.org_budget_id(org_id) + assert budget_id, "org created without a budget row" + deadline = time.monotonic() + 30 + while not any(row.budget_reset_at for row in client.budget_info(budget_id)): + if time.monotonic() > deadline: + pytest.fail("org budget window never scheduled by the reset job") + time.sleep(2) + + _drive_to_block(client, key) + _poll_until_serves_again(client, key) + + @pytest.mark.covers("quota_management.budget.internal_user.resets_after_window") + def test_personal_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)) + key = client.generate_key(user_id=user_id) + resources.defer(lambda: client.delete_key(key)) + + _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)