scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-02 07:13:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 148
litellm/tests/integration/management/test_team_budget_limits.py
devin-ai-integration[bot] 780c504e00
test: migrate tests/proxy_behavior into integration and unit tiers (#45939) 
* test: migrate tests/proxy_behavior into integration and unit tiers (partial)

* test(integration): run settings-dependent management tests on an owned proxy

* test(integration): assert stored values in the rewritten management rows

* test(integration): pin sibling team allocations against the org guaranteed throughput cap

* test(integration): keep Prisma clients out of the shared integration conftest import path

---------

Co-authored-by: yuneng <yuneng@berri.ai>
2026-10-11 08:28:42 -07:00

436 lines
15 KiB
Python
Raw Blame History

"""Phase 4 F3 — payload-level pins for team budget & rate-limit enforcement.
Pins the five helpers
* _check_team_model_specific_limits (team_endpoints.py:442)
* _check_team_rpm_tpm_limits (team_endpoints.py:527)
* check_org_team_model_specific_limits (team_endpoints.py:569)
* check_org_team_rpm_tpm_limits (team_endpoints.py:603)
* _check_org_team_limits (team_endpoints.py:628)
* _check_user_team_limits (team_endpoints.py:734)
Driven through /team/new + /team/update.
Structural finding, updated: /team/new and /team/update both load the org
via `get_org_object` WITH `include_budget_table=True`, so the org max_budget /
org tpm / org rpm guards inside `_check_org_team_limits` are live on both and
are pinned as enforced below. The `models` subset guard reads
`org_table.models` directly. The
`_check_user_team_limits` guards reach all branches through
`user_api_key_dict`, no relation include needed.
"""
import uuid
from typing import Any, Dict, Optional
import pytest
from litellm.proxy.utils import hash_token
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_org, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
async def _seed_scratch_actor_with_caps(
db,
scratch_prefix: str,
*,
user_role: str = "internal_user",
models: Optional[list] = None,
max_budget: Optional[float] = None,
tpm_limit: Optional[int] = None,
rpm_limit: Optional[int] = None,
) -> str:
"""Raw-seed a resources actor + verification token, with the token carrying
explicit caps that flow into `user_api_key_dict` at request time.
Returns cleartext key. Used by F3 user-team-limit scenarios where the
caller's caps drive the rejection. Uses 'internal_user' role plus a
token-level `allowed_routes` whitelist of /team/new + /team/update —
that whitelist is the only way past the admin-only role gate in
`RouteChecks.non_proxy_admin_allowed_routes_check`; without it the
non-admin caller would 401 before `_check_user_team_limits` ever fires.
"""
user_id = f"{scratch_prefix}-team-creator"
cleartext = "sk-" + uuid.uuid4().hex
await db.db.litellm_usertable.create(
data={
"user_id": user_id,
"user_role": user_role,
"max_budget": max_budget,
}
)
token_data: Dict[str, Any] = {
"token": hash_token(cleartext),
"key_name": f"{scratch_prefix}-team-creator-key",
"key_alias": f"{scratch_prefix}-team-creator-alias",
"user_id": user_id,
"models": models if models is not None else [],
"allowed_routes": ["/team/new", "/team/update"],
}
if tpm_limit is not None:
token_data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
token_data["rpm_limit"] = rpm_limit
await db.db.litellm_verificationtoken.create(data=token_data)
return cleartext
# ---------------------------------------------------------------------------
# _check_org_team_limits — models subset guard (the one path that reaches)
# ---------------------------------------------------------------------------
_ORG_MODEL_SCENARIOS = [
(
"org_models/team_subset_accepted",
["allowed-model"],
{"models": ["allowed-model"]},
200,
),
(
"org_models/team_extra_model_rejected",
["allowed-model"],
{"models": ["forbidden-model"]},
400,
),
(
"org_models/all_proxy_models_skips_check",
["all-proxy-models"],
{"models": ["any-model-at-all"]},
200,
),
]
@pytest.mark.parametrize(
"org_models,body_extras,expected_status",
[(b, c, d) for (_id, b, c, d) in _ORG_MODEL_SCENARIOS],
ids=[s[0] for s in _ORG_MODEL_SCENARIOS],
)
async def test_check_org_team_limits_models_subset(
org_models,
body_extras: Dict[str, Any],
expected_status: int,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(db, resources.prefix, models=org_models)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = resources.tag("team")
body: Dict[str, Any] = {
"team_id": team_id,
"team_alias": resources.prefix,
"organization_id": org_id,
**body_extras,
}
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
)
assert resp.status_code == expected_status, f"{body!r} → {resp.status_code}: {resp.text}"
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert len(rows) == (1 if expected_status == 200 else 0)
# ---------------------------------------------------------------------------
# _check_org_team_limits — budget / tpm / rpm live on /team/new and
# /team/update since both get_org_object calls pass include_budget_table=True.
# ---------------------------------------------------------------------------
_ORG_BUDGET_ENFORCED_SCENARIOS = [
(
"org_budget/over_max_budget_rejected",
{"max_budget": 100, "tpm_limit": None, "rpm_limit": None},
{"max_budget": 999_999},
400,
),
(
"org_budget/within_max_budget_accepted",
{"max_budget": 100, "tpm_limit": None, "rpm_limit": None},
{"max_budget": 50},
200,
),
(
"org_tpm/over_rejected",
{"max_budget": None, "tpm_limit": 100, "rpm_limit": None},
{"tpm_limit": 999_999},
400,
),
(
"org_tpm/within_accepted",
{"max_budget": None, "tpm_limit": 100, "rpm_limit": None},
{"tpm_limit": 50},
200,
),
(
"org_rpm/over_rejected",
{"max_budget": None, "tpm_limit": None, "rpm_limit": 100},
{"rpm_limit": 999_999},
400,
),
(
"org_rpm/within_accepted",
{"max_budget": None, "tpm_limit": None, "rpm_limit": 100},
{"rpm_limit": 50},
200,
),
]
@pytest.mark.parametrize(
"org_budget,body_extras,expected_status",
[(b, c, d) for (_id, b, c, d) in _ORG_BUDGET_ENFORCED_SCENARIOS],
ids=[s[0] for s in _ORG_BUDGET_ENFORCED_SCENARIOS],
)
async def test_check_org_team_limits_budget_enforced(
org_budget,
body_extras: Dict[str, Any],
expected_status: int,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(db, resources.prefix, **org_budget)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = resources.tag("team")
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json={
"team_id": team_id,
"team_alias": resources.prefix,
"organization_id": org_id,
**body_extras,
},
)
assert resp.status_code == expected_status, f"{body_extras!r} → {resp.status_code}: {resp.text}"
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert len(rows) == (1 if expected_status == 200 else 0)
@pytest.mark.parametrize(
"org_budget,body_extras,expected_status",
[(b, c, d) for (_id, b, c, d) in _ORG_BUDGET_ENFORCED_SCENARIOS],
ids=[s[0] for s in _ORG_BUDGET_ENFORCED_SCENARIOS],
)
async def test_check_org_team_limits_budget_enforced_on_update(
org_budget,
body_extras: Dict[str, Any],
expected_status: int,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(db, resources.prefix, **org_budget)
team_id = await create_scratch_team(db, resources.tag("team"), organization_id=org_id)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {seeder}"},
json={"team_id": team_id, **body_extras},
)
assert resp.status_code == expected_status, f"{body_extras!r} → {resp.status_code}: {resp.text}"
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
persisted = {field: getattr(row, field) for field in body_extras}
assert (persisted == body_extras) == (expected_status == 200)
# ---------------------------------------------------------------------------
# _check_user_team_limits — fires for standalone (no-org) teams created by
# a non-admin caller. Each guard reads from user_api_key_dict / user_obj.
# ---------------------------------------------------------------------------
_USER_LIMIT_SCENARIOS = [
(
"user_max_budget/within",
{"max_budget": 100.0, "models": [], "tpm_limit": None, "rpm_limit": None},
{"max_budget": 50.0},
200,
),
(
"user_max_budget/over",
{"max_budget": 100.0, "models": [], "tpm_limit": None, "rpm_limit": None},
{"max_budget": 1000.0},
400,
),
(
"user_models/subset",
{"max_budget": None, "models": ["m-a"], "tpm_limit": None, "rpm_limit": None},
{"models": ["m-a"]},
200,
),
(
"user_models/superset_rejected",
{"max_budget": None, "models": ["m-a"], "tpm_limit": None, "rpm_limit": None},
{"models": ["m-a", "m-b"]},
400,
),
(
"user_tpm/within",
{"max_budget": None, "models": [], "tpm_limit": 1000, "rpm_limit": None},
{"tpm_limit": 500},
200,
),
(
"user_tpm/over",
{"max_budget": None, "models": [], "tpm_limit": 1000, "rpm_limit": None},
{"tpm_limit": 2000},
400,
),
(
"user_rpm/within",
{"max_budget": None, "models": [], "tpm_limit": None, "rpm_limit": 100},
{"rpm_limit": 50},
200,
),
(
"user_rpm/over",
{"max_budget": None, "models": [], "tpm_limit": None, "rpm_limit": 100},
{"rpm_limit": 250},
400,
),
]
@pytest.mark.parametrize(
"actor_caps,body_extras,expected_status",
[(b, c, d) for (_id, b, c, d) in _USER_LIMIT_SCENARIOS],
ids=[s[0] for s in _USER_LIMIT_SCENARIOS],
)
async def test_check_user_team_limits(
actor_caps,
body_extras: Dict[str, Any],
expected_status: int,
gateway_client,
db,
resources,
):
caller = await _seed_scratch_actor_with_caps(db, resources.prefix, **actor_caps)
team_id = resources.tag("team")
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {caller}"},
json={
"team_id": team_id,
"team_alias": resources.prefix,
# Standalone team — no organization_id, so user-limit guard fires.
**body_extras,
},
)
assert resp.status_code == expected_status, (
f"caps={actor_caps} body={body_extras} → {resp.status_code}: {resp.text}"
)
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert len(rows) == (1 if expected_status == 200 else 0)
# ---------------------------------------------------------------------------
# /team/update path — budget authority.
#
# The caller's PERSONAL limits are never applied on update (that compared the
# wrong thing). Raising a team's spend ceiling is reserved for proxy admins.
# max_budget is not on the team-admin allow-list yet (LIT-5722), so a team
# admin is refused in either direction; the raise-only guard underneath the
# allow-list is pinned in the unit tests. _check_user_team_limits() only runs
# on /team/new.
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"personal_budget,requested_budget",
[(100000.0, 999.0), (10.0, 300.0)],
ids=["raise_with_generous_personal_budget", "lower_with_tiny_personal_budget"],
)
async def test_team_admin_cannot_change_budget_while_max_budget_is_not_editable(
gateway_client, db, resources, personal_budget: float, requested_budget: float
):
caller_cleartext = await _seed_scratch_actor_with_caps(db, resources.prefix, max_budget=personal_budget)
team_id = await create_scratch_team(
db,
team_id=resources.tag("team"),
admin_user_ids=[f"{resources.prefix}-team-creator"],
max_budget=500.0,
)
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {caller_cleartext}"},
json={"team_id": team_id, "max_budget": requested_budget},
)
assert resp.status_code == 403, resp.text
assert "Team admin editable fields" in resp.text, resp.text
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget == 500.0, "a refused update must leave the team budget unchanged"
async def test_proxy_admin_raise_budget_allowed(gateway_client, db, resources):
"""A proxy admin may raise a team's budget."""
team_id = await create_scratch_team(
db,
team_id=resources.tag("team"),
admin_user_ids=[f"{resources.prefix}-team-creator"],
max_budget=50.0,
)
# MASTER_KEY acts as proxy admin.
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
json={"team_id": team_id, "max_budget": 999.0},
)
assert resp.status_code == 200, resp.text
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget == 999.0, "proxy admin should be able to raise the budget"
async def test_team_admin_remove_budget_cap_blocked(gateway_client, db, resources):
"""A team admin cannot strip the team's cap (max_budget=null); removing the
ceiling is the strongest possible raise -> proxy-admin only."""
caller_cleartext = await _seed_scratch_actor_with_caps(db, resources.prefix, max_budget=100000.0)
team_id = await create_scratch_team(
db,
team_id=resources.tag("team"),
admin_user_ids=[f"{resources.prefix}-team-creator"],
max_budget=50.0,
)
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {caller_cleartext}"},
json={"team_id": team_id, "max_budget": None},
)
assert resp.status_code == 403, resp.text
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget == 50.0, "team budget cap must not be removed by a team admin"
async def test_proxy_admin_remove_budget_cap_allowed(gateway_client, db, resources):
"""A proxy admin may remove a team's cap (max_budget=null)."""
team_id = await create_scratch_team(
db,
team_id=resources.tag("team"),
admin_user_ids=[f"{resources.prefix}-team-creator"],
max_budget=50.0,
)
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
json={"team_id": team_id, "max_budget": None},
)
assert resp.status_code == 200, resp.text
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget is None, "proxy admin should be able to remove the cap"
Reference in a new issue View git blame Copy permalink