mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
200 lines
11 KiB
Python
200 lines
11 KiB
Python
from contextlib import ExitStack
|
|
from hashlib import sha256
|
|
from typing import Final
|
|
|
|
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 integration._support.client import Gateway, object_value
|
|
from integration._support.database import read_rows
|
|
from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
|
|
|
|
|
|
@pytest.mark.covers("mgmt.key.update.generated_sequences_preserve_state")
|
|
def test_generated_partial_updates_preserve_persisted_and_effective_state(gateway: Gateway) -> None:
|
|
class KeyUpdates(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.key = scenario.key(models=[self.models[0]], key_alias="initial", metadata={"revision": "initial"})
|
|
self.expected: dict[str, JsonValue] = {
|
|
"models": [self.models[0]], "key_alias": "initial", "metadata": {"revision": "initial"}
|
|
}
|
|
gateway.chat(self.models[0], key=self.key)
|
|
except BaseException:
|
|
with budget.cleanup():
|
|
self.resources.close()
|
|
raise
|
|
|
|
@rule(alias=st.sampled_from(("first", "second", "", "unicode-λ")))
|
|
def alias(self, alias: str) -> None:
|
|
gateway.post("/key/update", {"key": self.key, "key_alias": alias})
|
|
self.expected["key_alias"] = alias
|
|
|
|
@rule(index=st.integers(min_value=0, max_value=1), both=st.booleans())
|
|
def grant(self, index: int, both: bool) -> None:
|
|
models: Final = list(self.models) if both else [self.models[index]]
|
|
gateway.post("/key/update", {"key": self.key, "models": models})
|
|
self.expected["models"] = models
|
|
|
|
@rule(value=st.sampled_from(("", "a", "different", "λ")))
|
|
def metadata(self, value: str) -> None:
|
|
gateway.post("/key/update", {"key": self.key, "metadata": {"revision": value}})
|
|
self.expected["metadata"] = {"revision": value}
|
|
|
|
@invariant()
|
|
def persisted_state_and_serving_match(self) -> None:
|
|
rows: Final = read_rows(
|
|
'SELECT models, key_alias, metadata FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
|
(sha256(self.key.encode()).hexdigest(),),
|
|
)
|
|
assert rows == [self.expected]
|
|
info: Final = object_value(gateway.get("/key/info", {"key": self.key})["info"])
|
|
assert {field: info[field] for field in self.expected} == self.expected
|
|
for model in self.models:
|
|
response: Final = gateway.request(
|
|
"POST", "/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": "generated update control"}]},
|
|
key=self.key,
|
|
)
|
|
if model in self.expected["models"]:
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["usage"]["total_tokens"] == 40
|
|
else:
|
|
assert response.status_code == 403, response.text
|
|
assert response.json()["error"]["type"] == "key_model_access_denied"
|
|
|
|
def teardown(self) -> None:
|
|
with budget.cleanup():
|
|
self.resources.close()
|
|
|
|
with bounded_http_requests((gateway,), limit=2000) as budget:
|
|
run_state_machine_as_test(KeyUpdates, settings=LIFECYCLE_SETTINGS)
|
|
|
|
|
|
@pytest.mark.covers("mgmt.key.update.false_zero_and_empty_values_affect_serving")
|
|
def test_zero_false_and_empty_values_are_not_treated_as_omission(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
models: Final = (scenario.model(), scenario.model())
|
|
key: Final = scenario.key(models=[models[0]], max_budget=0, metadata={"ordinary": "value"})
|
|
denied: Final = gateway.request(
|
|
"POST", "/v1/chat/completions",
|
|
{"model": models[0], "messages": [{"role": "user", "content": "zero budget"}]}, key=key,
|
|
)
|
|
assert denied.status_code == 429, denied.text
|
|
assert denied.json()["error"]["type"] == "budget_exceeded"
|
|
gateway.post("/key/update", {"key": key, "max_budget": 1, "models": [], "metadata": {}})
|
|
info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
|
|
assert (info["models"], info["metadata"], info["max_budget"]) == ([], {}, 1)
|
|
for model in models:
|
|
assert object_value(gateway.chat(model, key=key)["usage"])["total_tokens"] == 40
|
|
gateway.post("/key/update", {"key": key, "blocked": True})
|
|
blocked: Final = gateway.request(
|
|
"POST", "/v1/chat/completions",
|
|
{"model": models[0], "messages": [{"role": "user", "content": "blocked control"}]}, key=key,
|
|
)
|
|
assert blocked.status_code == 401, blocked.text
|
|
assert blocked.json()["error"]["type"] == "auth_error"
|
|
gateway.post("/key/update", {"key": key, "blocked": False})
|
|
assert object_value(gateway.chat(models[0], key=key)["usage"])["total_tokens"] == 40
|
|
rows: Final = read_rows(
|
|
'SELECT blocked, models, metadata, max_budget FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
|
(sha256(key.encode()).hexdigest(),),
|
|
)
|
|
assert rows == [{"blocked": False, "models": [], "metadata": {}, "max_budget": 1.0}]
|
|
gateway.post("/key/update", {"key": key, "max_budget": 0})
|
|
assert read_rows(
|
|
'SELECT max_budget FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
|
(sha256(key.encode()).hexdigest(),),
|
|
) == [{"max_budget": 0.0}]
|
|
zero_after_update: Final = gateway.request(
|
|
"POST", "/v1/chat/completions",
|
|
{"model": models[0], "messages": [{"role": "user", "content": "updated zero budget"}]}, key=key,
|
|
)
|
|
assert zero_after_update.status_code == 429, zero_after_update.text
|
|
assert zero_after_update.json()["error"]["type"] == "budget_exceeded"
|
|
gateway.post("/key/update", {"key": key, "max_budget": None})
|
|
assert read_rows(
|
|
'SELECT max_budget FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
|
(sha256(key.encode()).hexdigest(),),
|
|
) == [{"max_budget": None}]
|
|
assert object_value(gateway.chat(models[0], key=key)["usage"])["total_tokens"] == 40
|
|
|
|
|
|
@pytest.mark.covers("mgmt.key.update.project_clear_preserves_scope", "mgmt.key.update.invalid_batch_is_atomic")
|
|
def test_project_omission_clear_and_invalid_update_have_distinct_effects(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model: Final = scenario.model()
|
|
outside: Final = scenario.model()
|
|
team: Final = scenario.team(models=[model])
|
|
project: Final = scenario.project(team, models=[model])
|
|
other: Final = scenario.project(team, models=[model])
|
|
key: Final = scenario.key(team_id=team, project_id=project, models=[model], key_alias="before", max_budget=5)
|
|
gateway.chat(model, key=key)
|
|
gateway.post("/key/update", {"key": key, "key_alias": "after"})
|
|
digest: Final = sha256(key.encode()).hexdigest()
|
|
|
|
def saved() -> list[dict[str, object]]:
|
|
return read_rows(
|
|
'SELECT key_alias, project_id, team_id, models, max_budget FROM "LiteLLM_VerificationToken" '
|
|
'WHERE token = %s', (digest,),
|
|
)
|
|
|
|
before: Final = saved()
|
|
assert before == [{"key_alias": "after", "project_id": project, "team_id": team, "models": [model], "max_budget": 5}]
|
|
for invalid in (other, ""):
|
|
rejected: Final = gateway.request(
|
|
"POST", "/key/update", {"key": key, "project_id": invalid, "key_alias": "must-not-persist"}
|
|
)
|
|
assert rejected.status_code == 400, rejected.text
|
|
assert saved() == before
|
|
gateway.chat(model, key=key)
|
|
gateway.post("/project/update", {"project_id": project, "blocked": True})
|
|
denied: Final = gateway.request(
|
|
"POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "blocked project"}]},
|
|
key=key,
|
|
)
|
|
assert denied.status_code == 401, denied.text
|
|
assert denied.json()["error"]["type"] == "auth_error"
|
|
for _ in range(2):
|
|
gateway.post("/key/update", {"key": key, "project_id": None})
|
|
assert saved() == [{**before[0], "project_id": None}]
|
|
assert object_value(gateway.chat(model, key=key)["usage"])["total_tokens"] == 40
|
|
outside_request: Final = gateway.request(
|
|
"POST", "/v1/chat/completions",
|
|
{"model": outside, "messages": [{"role": "user", "content": "detached scope control"}]}, key=key,
|
|
)
|
|
assert outside_request.status_code == 403, outside_request.text
|
|
assert outside_request.json()["error"]["type"] == "key_model_access_denied"
|
|
|
|
|
|
@pytest.mark.covers("mgmt.key.update.denied_request_preserves_effective_state")
|
|
def test_denied_key_update_preserves_saved_grants_and_serving(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model: Final = scenario.model()
|
|
outside: Final = scenario.model()
|
|
owner: Final = scenario.user(user_role="internal_user")
|
|
other: Final = scenario.user(user_role="internal_user")
|
|
key: Final = scenario.key(user_id=owner, models=[model], key_alias="unchanged", max_budget=2)
|
|
caller: Final = scenario.key(user_id=other, models=[model], allowed_routes=["/key/update", "/v1/chat/completions"])
|
|
gateway.chat(model, key=key)
|
|
denied: Final = gateway.request(
|
|
"POST", "/key/update", {"key": key, "key_alias": "wrong", "models": [outside], "max_budget": 0}, key=caller
|
|
)
|
|
assert denied.status_code == 403, denied.text
|
|
assert read_rows(
|
|
'SELECT user_id, models, key_alias, max_budget FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
|
(sha256(key.encode()).hexdigest(),),
|
|
) == [{"user_id": owner, "models": [model], "key_alias": "unchanged", "max_budget": 2.0}]
|
|
assert object_value(gateway.chat(model, key=key)["usage"])["total_tokens"] == 40
|
|
rejected: Final = gateway.request(
|
|
"POST", "/v1/chat/completions",
|
|
{"model": outside, "messages": [{"role": "user", "content": "unchanged scope"}]}, key=key,
|
|
)
|
|
assert rejected.status_code == 403, rejected.text
|
|
assert rejected.json()["error"]["type"] == "key_model_access_denied"
|