mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
780c504e00
* 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>
123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
import uuid
|
|
|
|
import pytest
|
|
|
|
from litellm.proxy.utils import hash_token
|
|
|
|
from tests.integration._support.actors import Actor
|
|
from tests.integration._support.actors import create_scratch_key
|
|
|
|
pytestmark = pytest.mark.asyncio(loop_scope="session")
|
|
|
|
_MARKER_BUDGET = 42.0
|
|
|
|
|
|
# POST /key/bulk_update is PROXY_ADMIN-only. The handler's own gate is
|
|
# user_role != PROXY_ADMIN -> 403, but ORG_ADMIN-role callers never reach it:
|
|
# the management-route gate 401s them first (the body carries no org context,
|
|
# and /key/bulk_update is an internal_user route, not an org-admin one).
|
|
# INTERNAL_USER-role callers clear the route gate and hit the handler's 403.
|
|
_MATRIX = [
|
|
("proxy_admin", Actor.PROXY_ADMIN, 200),
|
|
("org_admin", Actor.ORG_ADMIN, 401),
|
|
("team_admin", Actor.TEAM_ADMIN, 403),
|
|
("internal_user", Actor.INTERNAL_USER, 403),
|
|
("owner", Actor.OWNER, 403),
|
|
("unrelated_same_org", Actor.UNRELATED_SAME_ORG, 403),
|
|
("cross_org_user", Actor.CROSS_ORG_USER, 403),
|
|
("service_account", Actor.SERVICE_ACCOUNT, 403),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"actor,expected_status",
|
|
[(a, s) for (_id, a, s) in _MATRIX],
|
|
ids=[s[0] for s in _MATRIX],
|
|
)
|
|
async def test_key_bulk_update_authz_matrix(
|
|
actor: Actor,
|
|
expected_status: int,
|
|
gateway_client,
|
|
db,
|
|
resources,
|
|
actors,
|
|
):
|
|
caller = actors.keys[actor]
|
|
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
|
|
target = await create_scratch_key(
|
|
gateway_client, seeder, resources.prefix, user_id=caller.user_id
|
|
)
|
|
hashed = hash_token(target)
|
|
|
|
resp = await gateway_client.post(
|
|
"/key/bulk_update",
|
|
headers={"Authorization": f"Bearer {caller.cleartext}"},
|
|
json={"keys": [{"key": target, "max_budget": _MARKER_BUDGET}]},
|
|
)
|
|
assert (
|
|
resp.status_code == expected_status
|
|
), f"{actor.value}: {resp.status_code} {resp.text}"
|
|
|
|
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
|
|
assert row is not None
|
|
if expected_status == 200:
|
|
body = resp.json()
|
|
assert len(body["successful_updates"]) == 1
|
|
assert body["failed_updates"] == []
|
|
assert row.max_budget == _MARKER_BUDGET
|
|
else:
|
|
assert row.max_budget != _MARKER_BUDGET, "denied but key mutated"
|
|
|
|
|
|
async def test_key_bulk_update_empty_keys_is_400(gateway_client, actors):
|
|
"""An empty batch is rejected 400 before any per-key processing."""
|
|
resp = await gateway_client.post(
|
|
"/key/bulk_update",
|
|
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
|
|
json={"keys": []},
|
|
)
|
|
assert resp.status_code == 400, resp.text
|
|
|
|
|
|
async def test_key_bulk_update_over_max_batch_is_400(gateway_client, actors):
|
|
"""A batch larger than the 500-key cap is rejected 400."""
|
|
items = [{"key": "sk-" + uuid.uuid4().hex} for _ in range(501)]
|
|
resp = await gateway_client.post(
|
|
"/key/bulk_update",
|
|
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
|
|
json={"keys": items},
|
|
)
|
|
assert resp.status_code == 400, resp.text
|
|
|
|
|
|
async def test_key_bulk_update_per_key_failure_is_isolated(
|
|
gateway_client, db, resources, actors
|
|
):
|
|
"""One bad key in the batch does not abort the others — it lands in
|
|
failed_updates while the valid key is still updated."""
|
|
admin = actors.keys[Actor.PROXY_ADMIN]
|
|
valid = await create_scratch_key(
|
|
gateway_client, admin.cleartext, resources.prefix, user_id=admin.user_id
|
|
)
|
|
missing = "sk-" + uuid.uuid4().hex
|
|
|
|
resp = await gateway_client.post(
|
|
"/key/bulk_update",
|
|
headers={"Authorization": f"Bearer {admin.cleartext}"},
|
|
json={
|
|
"keys": [
|
|
{"key": valid, "max_budget": _MARKER_BUDGET},
|
|
{"key": missing, "max_budget": _MARKER_BUDGET},
|
|
]
|
|
},
|
|
)
|
|
assert resp.status_code == 200, resp.text
|
|
body = resp.json()
|
|
assert body["total_requested"] == 2
|
|
assert len(body["successful_updates"]) == 1
|
|
assert len(body["failed_updates"]) == 1
|
|
|
|
row = await db.db.litellm_verificationtoken.find_unique(
|
|
where={"token": hash_token(valid)}
|
|
)
|
|
assert row is not None and row.max_budget == _MARKER_BUDGET
|