mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-28 01:32:17 +00:00
41070b1363
* test(integration): pin the team-admin status-code matrix across every management door Every endpoint that admits a team admin today is called as a proxy admin, an admin of the target team, a plain member, an admin of another team and a teamless user, and the current status code is asserted per actor. The matrix is the parity check for collapsing the five team-admin helpers into one shared gate and for the later default-off permission flip. * test(integration): pin the permission-enabled team-admin doors in the gate matrix Adds three doors that run with team_admin_editable_team_fields granting max_budget, projects and member_key_budgets, so the enabled path is pinned alongside the default-off one. Hoists the ui_settings toggle from test_warmed_policy into the shared client so both files use one helper * test(integration): grant each permitted door only the permission it needs A door now names its single grant instead of every permission at once, so a gate that checks the wrong permission for a route turns that door red * test(integration): rewrite the team-admin matrix rows as request plus expected codes Each row now names the route it calls and the code each caller gets, and creates the member, key, model, callback or invitation it acts on through plain helpers on the shared team. Drops the Need, Target, World and Door types and the prepare step that seeded fixtures by enum.
269 lines
14 KiB
Python
269 lines
14 KiB
Python
import os
|
|
from contextlib import ExitStack
|
|
from hashlib import sha256
|
|
from typing import Final
|
|
|
|
import psycopg
|
|
import pytest
|
|
from hypothesis import strategies as st
|
|
from hypothesis.stateful import RuleBasedStateMachine, invariant, rule, run_state_machine_as_test
|
|
from pydantic import JsonValue
|
|
|
|
from tests.integration._support.client import Gateway, eventually, object_value, team_admin_permissions
|
|
from tests.integration._support.database import read_rows
|
|
from tests.integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
|
|
|
|
|
|
def assert_serving(gateway: Gateway, model: str, key: str, status: int, error_type: str = "auth_error") -> None:
|
|
response: Final = eventually(
|
|
lambda: gateway.request(
|
|
"POST", "/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": "warmed policy control"}]}, key=key,
|
|
),
|
|
lambda value: value.status_code == status,
|
|
seconds=3,
|
|
)
|
|
if status == 200:
|
|
assert response.json()["usage"]["total_tokens"] == 40
|
|
assert response.json()["choices"][0]["message"]["content"] == (
|
|
"Hello! This is a mock response from the fake OpenAI endpoint."
|
|
)
|
|
else:
|
|
assert response.json()["error"]["type"] == error_type
|
|
|
|
|
|
@pytest.mark.covers("mgmt.key.update.two_workers_enforce_warmed_policy")
|
|
def test_generated_policy_changes_reach_both_warmed_workers(gateway: Gateway, peer: Gateway) -> None:
|
|
class Policies(RuleBasedStateMachine):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.resources = ExitStack()
|
|
try:
|
|
scenario = self.resources.enter_context(gateway.scenario())
|
|
self.models = (scenario.model(), scenario.model())
|
|
self.allowed = 0
|
|
self.blocked = False
|
|
self.key = scenario.key(models=[self.models[0]], blocked=False)
|
|
self.control = scenario.key(models=list(self.models))
|
|
for worker in (gateway, peer):
|
|
assert_serving(worker, self.models[0], self.key, 200)
|
|
assert_serving(worker, self.models[1], self.control, 200)
|
|
except BaseException:
|
|
with budget.cleanup():
|
|
self.resources.close()
|
|
raise
|
|
|
|
@rule(index=st.integers(min_value=0, max_value=1))
|
|
def model_grant(self, index: int) -> None:
|
|
gateway.post("/key/update", {"key": self.key, "models": [self.models[index]]})
|
|
self.allowed = index
|
|
|
|
@rule(blocked=st.booleans())
|
|
def block(self, blocked: bool) -> None:
|
|
gateway.post("/key/update", {"key": self.key, "blocked": blocked})
|
|
self.blocked = blocked
|
|
|
|
@invariant()
|
|
def both_workers_enforce_policy(self) -> None:
|
|
rows: Final = read_rows(
|
|
'SELECT models, blocked FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
|
(sha256(self.key.encode()).hexdigest(),),
|
|
)
|
|
assert rows == [{"models": [self.models[self.allowed]], "blocked": self.blocked}]
|
|
for worker in (gateway, peer):
|
|
for index, model in enumerate(self.models):
|
|
status: Final = 401 if self.blocked else 200 if index == self.allowed else 403
|
|
kind: Final = "auth_error" if self.blocked else "key_model_access_denied"
|
|
assert_serving(worker, model, self.key, status, kind)
|
|
assert_serving(worker, self.models[1], self.control, 200)
|
|
|
|
def teardown(self) -> None:
|
|
with budget.cleanup():
|
|
self.resources.close()
|
|
|
|
with bounded_http_requests((gateway, peer), limit=3000) as budget:
|
|
run_state_machine_as_test(Policies, settings=LIFECYCLE_SETTINGS)
|
|
|
|
|
|
@pytest.mark.covers("mgmt.user.scim.deactivation_includes_nullable_blocked_keys")
|
|
def test_scim_deactivation_blocks_null_and_false_keys_but_preserves_other_owners(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model: Final = scenario.model()
|
|
user: Final = scenario.user(user_role="internal_user")
|
|
other: Final = scenario.user(user_role="internal_user")
|
|
null_key: Final = scenario.key(user_id=user, models=[model])
|
|
false_key: Final = scenario.key(user_id=user, models=[model], blocked=False)
|
|
manual: Final = scenario.key(user_id=user, models=[model], blocked=True)
|
|
control: Final = scenario.key(user_id=other, models=[model])
|
|
team: Final = scenario.team(models=[model])
|
|
service: Final = gateway.post("/key/service-account/generate", {"team_id": team, "models": [model]})
|
|
service_key: Final = service["key"]
|
|
assert isinstance(service_key, str)
|
|
scenario.cleanups.callback(scenario.delete_key, service_key)
|
|
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
|
|
connection.execute(
|
|
'UPDATE "LiteLLM_VerificationToken" SET blocked = NULL WHERE token = %s',
|
|
(sha256(null_key.encode()).hexdigest(),),
|
|
)
|
|
assert read_rows(
|
|
'SELECT user_id, blocked FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
|
(sha256(null_key.encode()).hexdigest(),),
|
|
) == [{"user_id": user, "blocked": None}]
|
|
assert read_rows(
|
|
'SELECT user_id FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
|
(sha256(service_key.encode()).hexdigest(),),
|
|
) == [{"user_id": None}]
|
|
for token in (null_key, false_key, control, service_key):
|
|
assert_serving(gateway, model, token, 200)
|
|
for active in (False, True):
|
|
response: Final = gateway.request(
|
|
"PATCH", f"/scim/v2/Users/{user}",
|
|
{"schemas": ["urn:ietf:params:scim:api:messages:2.0:PatchOp"],
|
|
"Operations": [{"op": "replace", "path": "active", "value": active}]},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
for token in (null_key, false_key):
|
|
rows: Final = read_rows(
|
|
'SELECT blocked, metadata FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
|
(sha256(token.encode()).hexdigest(),),
|
|
)
|
|
assert rows[0]["blocked"] is not active
|
|
assert object_value(rows[0]["metadata"]).get("scim_blocked") is (None if active else True)
|
|
assert_serving(gateway, model, token, 200 if active else 401)
|
|
assert_serving(gateway, model, manual, 401)
|
|
for token in (control, service_key):
|
|
assert_serving(gateway, model, token, 200)
|
|
|
|
|
|
@pytest.mark.covers("mgmt.team.member_update.demoted_role_cannot_write")
|
|
def test_warmed_team_role_demotion_prevents_later_management_writes(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario, team_admin_permissions(gateway, ["tpm_limit"]):
|
|
model: Final = scenario.model()
|
|
user: Final = scenario.user(user_role="internal_user")
|
|
team: Final = scenario.team(
|
|
models=[model], tpm_limit=1000, members_with_roles=[{"user_id": user, "role": "admin"}]
|
|
)
|
|
control_team: Final = scenario.team(models=[model], tpm_limit=1000)
|
|
caller: Final = scenario.key(
|
|
user_id=user, team_id=team, models=[model], allowed_routes=["/team/update", "/v1/chat/completions"]
|
|
)
|
|
gateway.chat(model, key=caller)
|
|
changed: Final = gateway.request("POST", "/team/update", {"team_id": team, "tpm_limit": 5000}, key=caller)
|
|
assert changed.status_code == 200, changed.text
|
|
assert read_rows('SELECT tpm_limit FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team,)) == [
|
|
{"tpm_limit": 5000}
|
|
]
|
|
unrelated_before: Final = read_rows(
|
|
'SELECT tpm_limit FROM "LiteLLM_TeamTable" WHERE team_id = %s', (control_team,)
|
|
)
|
|
unrelated: Final = gateway.request(
|
|
"POST", "/team/update", {"team_id": control_team, "tpm_limit": 7000}, key=caller
|
|
)
|
|
assert unrelated.status_code == 403, unrelated.text
|
|
assert read_rows(
|
|
'SELECT tpm_limit FROM "LiteLLM_TeamTable" WHERE team_id = %s', (control_team,)
|
|
) == unrelated_before
|
|
gateway.post("/team/member_update", {"team_id": team, "user_id": user, "role": "user"})
|
|
for target in (team, control_team):
|
|
before: Final = read_rows('SELECT tpm_limit FROM "LiteLLM_TeamTable" WHERE team_id = %s', (target,))
|
|
denied: Final = gateway.request(
|
|
"POST", "/team/update", {"team_id": target, "tpm_limit": 9000}, key=caller
|
|
)
|
|
assert denied.status_code == 403, denied.text
|
|
after: Final = read_rows('SELECT tpm_limit FROM "LiteLLM_TeamTable" WHERE team_id = %s', (target,))
|
|
assert after == before
|
|
roster: Final = read_rows('SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team,))
|
|
members: Final = roster[0]["members_with_roles"]
|
|
assert isinstance(members, list)
|
|
assert next(object_value(member)["role"] for member in members if object_value(member)["user_id"] == user) == "user"
|
|
assert_serving(gateway, model, caller, 200)
|
|
|
|
|
|
def _key_row(key: str) -> dict[str, JsonValue]:
|
|
rows: Final = read_rows(
|
|
'SELECT max_budget, key_alias FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
|
(sha256(key.encode()).hexdigest(),),
|
|
)
|
|
assert len(rows) == 1
|
|
return rows[0]
|
|
|
|
|
|
@pytest.mark.covers("mgmt.key.update.team_admin_member_key_budget_requires_opt_in")
|
|
def test_team_admin_changes_member_key_budget_only_when_opted_in(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model: Final = scenario.model()
|
|
admin: Final = scenario.user(user_role="internal_user")
|
|
member: Final = scenario.user(user_role="internal_user")
|
|
team: Final = scenario.team(
|
|
models=[model],
|
|
members_with_roles=[{"user_id": admin, "role": "admin"}, {"user_id": member, "role": "user"}],
|
|
)
|
|
other_team: Final = scenario.team(models=[model], members_with_roles=[{"user_id": member, "role": "user"}])
|
|
member_key: Final = scenario.key(
|
|
user_id=member, team_id=team, models=[model], max_budget=10, key_alias="member"
|
|
)
|
|
personal_key: Final = scenario.key(user_id=member, models=[model], max_budget=10)
|
|
foreign_key: Final = scenario.key(user_id=member, team_id=other_team, models=[model], max_budget=10)
|
|
admin_key: Final = scenario.key(
|
|
user_id=admin, team_id=team, models=[model], allowed_routes=["/key/update", "/v1/chat/completions"]
|
|
)
|
|
member_caller: Final = scenario.key(
|
|
user_id=member, team_id=team, models=[model], allowed_routes=["/key/update", "/v1/chat/completions"]
|
|
)
|
|
assert_serving(gateway, model, member_key, 200)
|
|
with team_admin_permissions(gateway, []):
|
|
denied: Final = gateway.request("POST", "/key/update", {"key": member_key, "max_budget": 0}, key=admin_key)
|
|
assert denied.status_code == 403, denied.text
|
|
assert _key_row(member_key) == {"max_budget": 10.0, "key_alias": "member"}
|
|
with team_admin_permissions(gateway, ["member_key_budgets"]):
|
|
for target in (personal_key, foreign_key):
|
|
out_of_scope: Final = gateway.request(
|
|
"POST", "/key/update", {"key": target, "max_budget": 0}, key=admin_key
|
|
)
|
|
assert out_of_scope.status_code == 403, out_of_scope.text
|
|
assert _key_row(target)["max_budget"] == 10.0
|
|
by_member: Final = gateway.request(
|
|
"POST", "/key/update", {"key": admin_key, "max_budget": 0}, key=member_caller
|
|
)
|
|
assert by_member.status_code == 403, by_member.text
|
|
not_budget: Final = gateway.request(
|
|
"POST", "/key/update", {"key": member_key, "key_alias": "renamed"}, key=admin_key
|
|
)
|
|
assert not_budget.status_code == 403, not_budget.text
|
|
assert _key_row(member_key) == {"max_budget": 10.0, "key_alias": "member"}
|
|
changed: Final = gateway.request(
|
|
"POST", "/key/update", {"key": member_key, "max_budget": 0, "budget_duration": "30d"}, key=admin_key
|
|
)
|
|
assert changed.status_code == 200, changed.text
|
|
assert _key_row(member_key) == {"max_budget": 0.0, "key_alias": "member"}
|
|
assert_serving(gateway, model, member_key, 422, "budget_exceeded")
|
|
restored: Final = gateway.request(
|
|
"POST", "/key/update", {"key": member_key, "max_budget": 10}, key=admin_key
|
|
)
|
|
assert restored.status_code == 200, restored.text
|
|
assert_serving(gateway, model, member_key, 200)
|
|
|
|
|
|
@pytest.mark.covers("mgmt.key.update.expiry_changes_reach_warmed_workers")
|
|
def test_expiry_and_explicit_clear_reach_both_warmed_workers(gateway: Gateway, peer: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model: Final = scenario.model()
|
|
key: Final = scenario.key(models=[model], duration="1h")
|
|
control: Final = scenario.key(models=[model])
|
|
for worker in (gateway, peer):
|
|
assert_serving(worker, model, key, 200)
|
|
gateway.post("/key/update", {"key": key, "duration": "0s"})
|
|
assert read_rows(
|
|
"SELECT expires <= timezone('UTC', now()) AS expired FROM \"LiteLLM_VerificationToken\" WHERE token = %s",
|
|
(sha256(key.encode()).hexdigest(),),
|
|
) == [{"expired": True}]
|
|
for worker in (gateway, peer):
|
|
assert_serving(worker, model, key, 401, "expired_key")
|
|
assert_serving(worker, model, control, 200)
|
|
gateway.post("/key/update", {"key": key, "duration": None})
|
|
assert read_rows(
|
|
'SELECT expires IS NULL AS cleared FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
|
(sha256(key.encode()).hexdigest(),),
|
|
) == [{"cleared": True}]
|
|
for worker in (gateway, peer):
|
|
assert_serving(worker, model, key, 200)
|