mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
44d99475a5
* test(integration): stop GET /__observations from draining the upstream log GET /__observations used to empty the queue, so a second read in a test saw nothing and a "never reached upstream" check could pass whether or not the request was forwarded. GET now returns every request since the last reset and DELETE /__observations is the reset. An autouse fixture resets the worker upstream before each test, which also covers tests that read without the gateway fixture. Discarded reads used as resets become DELETEs, helpers that accumulated drained reads now read one snapshot, and tests with several phases reset between them. The three vacuous negative checks (rag allowlist, guardrail opt-out chaos, xAI batch output) now assert the absence on the same read that shows a request which did reach the upstream. * style(integration): space the observations unit test like its neighbours
261 lines
12 KiB
Python
261 lines
12 KiB
Python
import uuid
|
|
from collections.abc import Iterator
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
from integration._support.client import Gateway, Scenario, object_value, string_value
|
|
from integration._support.database import read_rows
|
|
from pydantic import BaseModel, JsonValue
|
|
|
|
|
|
class _ObservedRequest(BaseModel):
|
|
body: dict[str, JsonValue]
|
|
|
|
|
|
class _Observations(BaseModel):
|
|
requests: list[_ObservedRequest]
|
|
|
|
|
|
class _TeamState(BaseModel):
|
|
team_id: str
|
|
models: list[str] | None = None
|
|
|
|
|
|
class _TeamModelRow(BaseModel):
|
|
models: list[str] | None = None
|
|
|
|
|
|
class _TeamInfo(BaseModel):
|
|
team_info: _TeamState
|
|
|
|
|
|
class _ProxyError(BaseModel):
|
|
type: str
|
|
|
|
|
|
class _ErrorResponse(BaseModel):
|
|
error: _ProxyError
|
|
|
|
|
|
@pytest.fixture
|
|
def upstream(gateway: Gateway) -> Iterator[httpx.Client]:
|
|
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as client:
|
|
client.delete("/__observations").raise_for_status()
|
|
yield client
|
|
|
|
|
|
def _observed_requests(upstream: httpx.Client) -> list[JsonValue]:
|
|
observed: Final = upstream.get("/__observations")
|
|
observed.raise_for_status()
|
|
requests: Final = object_value(observed.json())["requests"]
|
|
assert isinstance(requests, list)
|
|
return requests
|
|
|
|
|
|
def _calls_to(observed: list[JsonValue], provider_model: str) -> int:
|
|
return sum(object_value(object_value(request)["body"]).get("model") == provider_model for request in observed)
|
|
|
|
|
|
def _observed_provider_models(upstream: httpx.Client) -> tuple[str, ...]:
|
|
observed: Final = upstream.get("/__observations")
|
|
assert observed.status_code == 200, observed.text
|
|
response: Final = _Observations.model_validate_json(observed.text)
|
|
return tuple(string_value(request.body["model"]) for request in response.requests)
|
|
|
|
|
|
def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response:
|
|
return gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": f"team model {uuid.uuid4().hex}"}]},
|
|
key=key,
|
|
)
|
|
|
|
|
|
def _team_model(scenario: Scenario) -> tuple[str, str]:
|
|
provider_model: Final = f"team-model-{uuid.uuid4().hex}"
|
|
return scenario.model(model=f"openai/{provider_model}"), provider_model
|
|
|
|
|
|
def _team_info(gateway: Gateway, team_id: str) -> _TeamInfo:
|
|
response: Final = gateway.request("GET", "/team/info", params={"team_id": team_id})
|
|
assert response.status_code == 200, response.text
|
|
return _TeamInfo.model_validate_json(response.text)
|
|
|
|
|
|
def _stored_team_models(team_id: str) -> _TeamModelRow:
|
|
rows: Final = read_rows('SELECT models FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,))
|
|
assert len(rows) == 1, repr(rows)
|
|
return _TeamModelRow.model_validate(rows[0])
|
|
|
|
|
|
def _model_denied(response: httpx.Response) -> None:
|
|
assert response.status_code == 403, response.text
|
|
assert _ErrorResponse.model_validate_json(response.text).error.type == "team_model_access_denied", response.text
|
|
|
|
|
|
def _ids(listing: dict[str, JsonValue]) -> set[JsonValue]:
|
|
data: Final = listing["data"]
|
|
assert isinstance(data, list)
|
|
return {object_value(entry)["id"] for entry in data}
|
|
|
|
|
|
def test_a_model_created_for_a_team_is_listed_in_that_teams_models(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
team: Final = scenario.team(models=[])
|
|
other: Final = scenario.team(models=[])
|
|
model: Final = scenario.model(model_info={"team_id": team})
|
|
own_models: Final = object_value(gateway.get("/team/info", {"team_id": team})["team_info"])["models"]
|
|
other_models: Final = object_value(gateway.get("/team/info", {"team_id": other})["team_info"])["models"]
|
|
assert isinstance(own_models, list) and isinstance(other_models, list)
|
|
assert model in own_models
|
|
assert model not in other_models
|
|
|
|
|
|
def test_a_team_model_is_listed_and_served_only_for_keys_of_its_team(gateway: Gateway, upstream: httpx.Client) -> None:
|
|
with gateway.scenario() as scenario:
|
|
team: Final = scenario.team(models=[])
|
|
other: Final = scenario.team(models=[])
|
|
provider_model: Final = f"team-scoped-{uuid.uuid4().hex}"
|
|
model: Final = scenario.model(model=f"openai/{provider_model}", model_info={"team_id": team})
|
|
team_key: Final = scenario.key(team_id=team)
|
|
other_key: Final = scenario.key(team_id=other)
|
|
assert model in _ids(object_value(gateway.request("GET", "/models", key=team_key).json()))
|
|
assert model not in _ids(object_value(gateway.request("GET", "/models", key=other_key).json()))
|
|
served: Final = _chat(gateway, model, team_key)
|
|
assert served.status_code == 200, served.text
|
|
refused: Final = _chat(gateway, model, other_key)
|
|
assert refused.status_code == 400, refused.text
|
|
observed: Final = _observed_requests(upstream)
|
|
assert _calls_to(observed, provider_model) == 1, observed
|
|
|
|
|
|
def _v2_team_public_names(gateway: Gateway, key: str, model: str) -> list[JsonValue]:
|
|
response: Final = gateway.request("GET", "/v2/model/info", key=key, params={"model": model})
|
|
assert response.status_code == 200, response.text
|
|
return [entry["model_info"].get("team_public_model_name") for entry in response.json()["data"]]
|
|
|
|
|
|
def test_v2_model_info_reports_a_team_model_to_team_and_non_team_keys(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
team: Final = scenario.team(models=[])
|
|
model: Final = scenario.model(model_info={"team_id": team})
|
|
assert _v2_team_public_names(gateway, scenario.key(team_id=team), model) == [model]
|
|
assert _v2_team_public_names(gateway, scenario.key(), model) == [model]
|
|
|
|
|
|
@pytest.mark.parametrize("set_on", ["team_new", "team_update"])
|
|
def test_team_model_alias_routes_a_team_key_to_its_target(
|
|
gateway: Gateway, upstream: httpx.Client, set_on: str
|
|
) -> None:
|
|
with gateway.scenario() as scenario:
|
|
provider_model: Final = f"team-alias-{uuid.uuid4().hex}"
|
|
model: Final = scenario.model(model=f"openai/{provider_model}")
|
|
alias: Final = f"alias-{uuid.uuid4().hex}"
|
|
team: Final = (
|
|
scenario.team(models=[model], model_aliases={alias: model})
|
|
if set_on == "team_new"
|
|
else scenario.team(models=[model])
|
|
)
|
|
if set_on == "team_update":
|
|
gateway.post("/team/update", {"team_id": team, "model_aliases": {alias: model}})
|
|
key: Final = scenario.key(team_id=team, models=[model])
|
|
response: Final = _chat(gateway, alias, key)
|
|
assert response.status_code == 200, response.text
|
|
assert string_value(response.json()["model"]) == alias
|
|
observed: Final = _observed_requests(upstream)
|
|
assert _calls_to(observed, provider_model) == 1, observed
|
|
upstream.delete("/__observations").raise_for_status()
|
|
unaliased: Final = _chat(gateway, f"alias-{uuid.uuid4().hex}", key)
|
|
assert unaliased.status_code == 403, unaliased.text
|
|
assert unaliased.json()["error"]["type"] == "key_model_access_denied"
|
|
after_refusal: Final = _observed_requests(upstream)
|
|
assert _calls_to(after_refusal, provider_model) == 0, after_refusal
|
|
|
|
|
|
def test_team_model_add_and_delete_change_what_a_warmed_team_key_can_call(
|
|
gateway: Gateway, upstream: httpx.Client
|
|
) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model_one, provider_one = _team_model(scenario)
|
|
model_two, provider_two = _team_model(scenario)
|
|
model_three, provider_three = _team_model(scenario)
|
|
team: Final = scenario.team(models=[model_one])
|
|
key: Final = scenario.key(team_id=team)
|
|
|
|
first: Final = _chat(gateway, model_one, key)
|
|
assert first.status_code == 200, first.text
|
|
initially_denied: Final = _chat(gateway, model_two, key)
|
|
_model_denied(initially_denied)
|
|
observed_before_add: Final = _observed_provider_models(upstream)
|
|
assert observed_before_add == (provider_one,), repr(observed_before_add)
|
|
upstream.delete("/__observations").raise_for_status()
|
|
|
|
added: Final = gateway.request("POST", "/team/model/add", {"team_id": team, "models": [model_two, model_three]})
|
|
assert added.status_code == 200, added.text
|
|
added_state: Final = _TeamState.model_validate_json(added.text)
|
|
expected_added: Final = sorted((model_one, model_two, model_three))
|
|
assert sorted(added_state.models or ()) == expected_added, added.text
|
|
added_info: Final = _team_info(gateway, team)
|
|
assert sorted(added_info.team_info.models or ()) == expected_added, repr(added_info)
|
|
added_row: Final = _stored_team_models(team)
|
|
assert sorted(added_row.models or ()) == expected_added, repr(added_row)
|
|
|
|
second: Final = _chat(gateway, model_two, key)
|
|
assert second.status_code == 200, second.text
|
|
third: Final = _chat(gateway, model_three, key)
|
|
assert third.status_code == 200, third.text
|
|
deleted: Final = gateway.request("POST", "/team/model/delete", {"team_id": team, "models": [model_two]})
|
|
assert deleted.status_code == 200, deleted.text
|
|
deleted_state: Final = _TeamState.model_validate_json(deleted.text)
|
|
expected_remaining: Final = sorted((model_one, model_three))
|
|
assert sorted(deleted_state.models or ()) == expected_remaining, deleted.text
|
|
deleted_info: Final = _team_info(gateway, team)
|
|
assert sorted(deleted_info.team_info.models or ()) == expected_remaining, repr(deleted_info)
|
|
deleted_row: Final = _stored_team_models(team)
|
|
assert sorted(deleted_row.models or ()) == expected_remaining, repr(deleted_row)
|
|
|
|
denied_after_delete: Final = _chat(gateway, model_two, key)
|
|
_model_denied(denied_after_delete)
|
|
still_allowed_one: Final = _chat(gateway, model_one, key)
|
|
assert still_allowed_one.status_code == 200, still_allowed_one.text
|
|
still_allowed_three: Final = _chat(gateway, model_three, key)
|
|
assert still_allowed_three.status_code == 200, still_allowed_three.text
|
|
observed_after_add_and_delete: Final = _observed_provider_models(upstream)
|
|
assert observed_before_add + observed_after_add_and_delete == (
|
|
provider_one,
|
|
provider_two,
|
|
provider_three,
|
|
provider_one,
|
|
provider_three,
|
|
), repr((observed_before_add, observed_after_add_and_delete))
|
|
|
|
|
|
def test_team_model_add_on_an_unrestricted_team_keeps_every_other_model(
|
|
gateway: Gateway, upstream: httpx.Client
|
|
) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model_one, provider_one = _team_model(scenario)
|
|
model_two, provider_two = _team_model(scenario)
|
|
team: Final = scenario.team(models=[])
|
|
key: Final = scenario.key(team_id=team)
|
|
|
|
before_add: Final = _chat(gateway, model_two, key)
|
|
assert before_add.status_code == 200, before_add.text
|
|
|
|
added: Final = gateway.request("POST", "/team/model/add", {"team_id": team, "models": [model_one]})
|
|
assert added.status_code == 200, added.text
|
|
expected: Final = sorted(("all-proxy-models", model_one))
|
|
added_state: Final = _TeamState.model_validate_json(added.text)
|
|
assert sorted(added_state.models or ()) == expected, added.text
|
|
added_info: Final = _team_info(gateway, team)
|
|
assert sorted(added_info.team_info.models or ()) == expected, repr(added_info)
|
|
added_row: Final = _stored_team_models(team)
|
|
assert sorted(added_row.models or ()) == expected, repr(added_row)
|
|
still_unrestricted: Final = _chat(gateway, model_two, key)
|
|
assert still_unrestricted.status_code == 200, still_unrestricted.text
|
|
added_model: Final = _chat(gateway, model_one, key)
|
|
assert added_model.status_code == 200, added_model.text
|
|
observed_after_add: Final = _observed_provider_models(upstream)
|
|
assert observed_after_add == (provider_two, provider_two, provider_one), repr(observed_after_add)
|