scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

test(e2e): cover key max_budget blocks on personal, team, and team-member keys (#33895)

Browse source 
* 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
This commit is contained in:
ryan-crabbe-berri 2026-07-22 17:48:18 -07:00 • committed by GitHub
parent 301a02b7be
commit 9f6b3d24f7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 147 additions and 250 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
tests/e2e/conftest.py
View file

@ -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.

42
tests/e2e/lifecycle.py
View file

@ -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

349
tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py
View file

@ -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))