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
372 lines
17 KiB
Python
372 lines
17 KiB
Python
import uuid
|
|
from contextlib import ExitStack
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
import yaml
|
|
from integration._support.client import Gateway, eventually, object_value
|
|
from integration._support.database import read_rows
|
|
from integration._support.process import owned_proxy
|
|
from pydantic import BaseModel, JsonValue
|
|
|
|
pytestmark: Final = pytest.mark.shares_worker_state
|
|
|
|
|
|
class _BlockedCustomerInfo(BaseModel):
|
|
user_id: str
|
|
blocked: bool
|
|
|
|
|
|
class _BlockUsersResponse(BaseModel):
|
|
blocked_users: list[_BlockedCustomerInfo]
|
|
|
|
|
|
def _blocked_row(user_id: str) -> list[dict[str, JsonValue]]:
|
|
return read_rows('SELECT blocked FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_id,))
|
|
|
|
|
|
def _observed(upstream: httpx.Client) -> list[dict[str, JsonValue]]:
|
|
response: Final = upstream.get("/__observations")
|
|
response.raise_for_status()
|
|
requests: Final = response.json()["requests"]
|
|
assert isinstance(requests, list)
|
|
return requests
|
|
|
|
|
|
def _reset_observations(upstream: httpx.Client) -> None:
|
|
upstream.delete("/__observations").raise_for_status()
|
|
|
|
|
|
def _chat(gateway: Gateway, model: str, key: str, end_user: str) -> httpx.Response:
|
|
return gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": f"end user block {uuid.uuid4().hex}"}],
|
|
"user": end_user,
|
|
},
|
|
key=key,
|
|
)
|
|
|
|
|
|
def _observation(text: str, end_user: str) -> dict[str, JsonValue]:
|
|
return {
|
|
"path": "/v1/chat/completions",
|
|
"authorization": "Bearer integration-provider-key",
|
|
"body": {
|
|
"messages": [{"role": "user", "content": text}],
|
|
"model": "gpt-4o-mini",
|
|
"user": end_user,
|
|
},
|
|
"method": "POST",
|
|
"api_key": "",
|
|
}
|
|
|
|
|
|
def _assert_blocked_refusal(response: httpx.Response, end_user: str) -> None:
|
|
assert response.status_code == 400, response.text
|
|
assert response.json() == {
|
|
"error": {
|
|
"message": f"User blocked from making LLM API Calls. User={end_user}",
|
|
"type": "invalid_request_error",
|
|
"param": None,
|
|
"code": "400",
|
|
"provider_specific_fields": {"error": f"User blocked from making LLM API Calls. User={end_user}"},
|
|
}
|
|
}, response.text
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _CallbackRig:
|
|
first: Gateway
|
|
second: Gateway
|
|
|
|
|
|
def _callback_config(directory: Path, blocked_user_list: list[str] | None = None) -> Path:
|
|
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
|
settings: Final = {
|
|
**object_value(config["litellm_settings"]),
|
|
"callbacks": ["blocked_user_check"],
|
|
}
|
|
if blocked_user_list is not None:
|
|
settings["blocked_user_list"] = blocked_user_list
|
|
config["litellm_settings"] = settings
|
|
path: Final = directory / "config.yaml"
|
|
path.write_text(yaml.safe_dump(config))
|
|
return path
|
|
|
|
|
|
def _two_callback_proxies(gateway: Gateway, tmp_path: Path, stack: ExitStack) -> _CallbackRig:
|
|
config: Final = _callback_config(tmp_path)
|
|
first: Final = stack.enter_context(owned_proxy(gateway, tmp_path / "a", {}, config=config))
|
|
second: Final = stack.enter_context(owned_proxy(gateway, tmp_path / "b", {}, config=config))
|
|
return _CallbackRig(first, second)
|
|
|
|
|
|
def _one_callback_proxy(gateway: Gateway, tmp_path: Path, stack: ExitStack) -> Gateway:
|
|
config: Final = _callback_config(tmp_path)
|
|
return stack.enter_context(owned_proxy(gateway, tmp_path / "a", {}, config=config))
|
|
|
|
|
|
@pytest.mark.parametrize("block_route", ("/customer/block", "/end_user/block"))
|
|
def test_block_route_marks_the_customer_blocked_and_update_unblocks(gateway: Gateway, block_route: str) -> None:
|
|
with gateway.scenario() as scenario:
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
created: Final = gateway.request("POST", "/customer/new", {"user_id": end_user, "max_budget": 100.0})
|
|
assert created.status_code == 200, created.text
|
|
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
|
|
blocked: Final = gateway.request("POST", block_route, {"user_ids": [end_user]})
|
|
assert blocked.status_code == 200, blocked.text
|
|
assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": True}], seconds=15) == [
|
|
{"blocked": True}
|
|
]
|
|
info: Final = gateway.get("/customer/info", {"end_user_id": end_user})
|
|
assert _BlockedCustomerInfo.model_validate(info).blocked is True
|
|
|
|
updated: Final = gateway.request("POST", "/customer/update", {"user_id": end_user, "blocked": False})
|
|
assert updated.status_code == 200, updated.text
|
|
assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": False}], seconds=15) == [
|
|
{"blocked": False}
|
|
]
|
|
assert (
|
|
_BlockedCustomerInfo.model_validate(gateway.get("/customer/info", {"end_user_id": end_user})).blocked
|
|
is False
|
|
)
|
|
|
|
|
|
def _assert_unseen_customer_is_blocked(gateway: Gateway, block_route: str) -> None:
|
|
with gateway.scenario() as scenario:
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
|
|
blocked: Final = gateway.request("POST", block_route, {"user_ids": [end_user]})
|
|
assert blocked.status_code == 200, blocked.text
|
|
blocked_response: Final = _BlockUsersResponse.model_validate_json(blocked.content)
|
|
assert blocked_response == _BlockUsersResponse(
|
|
blocked_users=[_BlockedCustomerInfo(user_id=end_user, blocked=True)]
|
|
)
|
|
assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": True}], seconds=15) == [
|
|
{"blocked": True}
|
|
]
|
|
info: Final = gateway.request("GET", "/customer/info", params={"end_user_id": end_user})
|
|
assert info.status_code == 200, info.text
|
|
assert _BlockedCustomerInfo.model_validate(info.json()).blocked is True
|
|
|
|
updated: Final = gateway.request("POST", "/customer/update", {"user_id": end_user, "blocked": False})
|
|
assert updated.status_code == 200, updated.text
|
|
assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": False}], seconds=15) == [
|
|
{"blocked": False}
|
|
]
|
|
info_after_update: Final = gateway.request("GET", "/customer/info", params={"end_user_id": end_user})
|
|
assert info_after_update.status_code == 200, info_after_update.text
|
|
assert _BlockedCustomerInfo.model_validate(info_after_update.json()).blocked is False
|
|
|
|
|
|
def test_end_user_block_route_creates_and_blocks_an_unseen_customer(gateway: Gateway) -> None:
|
|
_assert_unseen_customer_is_blocked(gateway, "/end_user/block")
|
|
|
|
|
|
def test_customer_block_route_creates_and_blocks_an_unseen_customer(gateway: Gateway) -> None:
|
|
_assert_unseen_customer_is_blocked(gateway, "/customer/block")
|
|
|
|
|
|
def test_listed_end_user_is_refused_and_unblock_restores_serving(gateway: Gateway, tmp_path: Path) -> None:
|
|
with (
|
|
ExitStack() as stack,
|
|
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
|
):
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
config: Final = _callback_config(tmp_path, blocked_user_list=[end_user])
|
|
proxy: Final = stack.enter_context(owned_proxy(gateway, tmp_path / "listed", {}, config=config))
|
|
with proxy.scenario() as scenario:
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
key: Final = scenario.key(models=[model])
|
|
assert _observed(upstream) == []
|
|
|
|
refused: Final = _chat(proxy, model, key, end_user)
|
|
_assert_blocked_refusal(refused, end_user)
|
|
assert _observed(upstream) == []
|
|
|
|
unblocked: Final = proxy.request("POST", "/customer/unblock", {"user_ids": [end_user]})
|
|
assert unblocked.status_code == 200, unblocked.text
|
|
assert unblocked.json() == {"blocked_users": []}, unblocked.text
|
|
assert _blocked_row(end_user) == []
|
|
|
|
text: Final = f"end user unblocked {uuid.uuid4().hex}"
|
|
served: Final = proxy.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": text}], "user": end_user},
|
|
key=key,
|
|
)
|
|
assert served.status_code == 200, served.text
|
|
assert _observed(upstream) == [_observation(text, end_user)]
|
|
|
|
|
|
@pytest.mark.parametrize("block_route", ("/customer/block", "/end_user/block"))
|
|
def test_blocked_end_user_is_refused_on_every_proxy_and_unblock_restores(
|
|
gateway: Gateway, tmp_path: Path, block_route: str
|
|
) -> None:
|
|
pytest.skip("BUG: blocked end-user chat requests return HTTP 200 instead of HTTP 400")
|
|
with (
|
|
ExitStack() as stack,
|
|
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
|
):
|
|
rig: Final = _two_callback_proxies(gateway, tmp_path, stack)
|
|
with rig.first.scenario() as scenario:
|
|
_reset_observations(upstream)
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
created: Final = rig.first.request("POST", "/customer/new", {"user_id": end_user, "max_budget": 100.0})
|
|
assert created.status_code == 200, created.text
|
|
scenario.cleanups.callback(rig.first.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
key: Final = scenario.key(models=[model])
|
|
served: Final = _chat(rig.first, model, key, end_user)
|
|
assert served.status_code == 200, served.text
|
|
_reset_observations(upstream)
|
|
|
|
rig.first.request("POST", block_route, {"user_ids": [end_user]})
|
|
assert eventually(
|
|
lambda: _blocked_row(end_user), lambda value: value == [{"blocked": True}], seconds=15
|
|
) == [{"blocked": True}]
|
|
info: Final = rig.first.get("/customer/info", {"end_user_id": end_user})
|
|
assert _BlockedCustomerInfo.model_validate(info).blocked is True
|
|
|
|
denied_first: Final = _chat(rig.first, model, key, end_user)
|
|
denied_second: Final = _chat(rig.second, model, key, end_user)
|
|
for denied in (denied_first, denied_second):
|
|
_assert_blocked_refusal(denied, end_user)
|
|
upstream_calls: Final = _observed(upstream)
|
|
assert upstream_calls == [], upstream_calls
|
|
|
|
updated: Final = rig.first.request("POST", "/customer/update", {"user_id": end_user, "blocked": False})
|
|
assert updated.status_code == 200, updated.text
|
|
assert eventually(
|
|
lambda: _blocked_row(end_user), lambda value: value == [{"blocked": False}], seconds=15
|
|
) == [{"blocked": False}]
|
|
assert (
|
|
_BlockedCustomerInfo.model_validate(rig.first.get("/customer/info", {"end_user_id": end_user})).blocked
|
|
is False
|
|
)
|
|
|
|
text: Final = f"end user unblocked {uuid.uuid4().hex}"
|
|
|
|
def call(candidate: Gateway) -> httpx.Response:
|
|
return candidate.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": text}], "user": end_user},
|
|
key=key,
|
|
)
|
|
|
|
assert (
|
|
eventually(
|
|
lambda: call(rig.first), lambda response: response.status_code == 200, seconds=10
|
|
).status_code
|
|
== 200
|
|
)
|
|
assert (
|
|
eventually(
|
|
lambda: call(rig.second), lambda response: response.status_code == 200, seconds=10
|
|
).status_code
|
|
== 200
|
|
)
|
|
expected_observations: Final = [
|
|
{
|
|
"path": "/v1/chat/completions",
|
|
"authorization": "Bearer integration-provider-key",
|
|
"body": {
|
|
"messages": [{"role": "user", "content": text}],
|
|
"model": "gpt-4o-mini",
|
|
"user": end_user,
|
|
},
|
|
"method": "POST",
|
|
"api_key": "",
|
|
}
|
|
] * 2
|
|
assert _observed(upstream) == expected_observations
|
|
|
|
|
|
def test_unblock_route_requires_the_blocked_user_check_callback(gateway: Gateway, tmp_path: Path) -> None:
|
|
with ExitStack() as stack:
|
|
callback_proxy: Final = _one_callback_proxy(gateway, tmp_path, stack)
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
created: Final = callback_proxy.request(
|
|
"POST", "/customer/new", {"user_id": end_user, "max_budget": 100.0, "blocked": True}
|
|
)
|
|
assert created.status_code == 200, created.text
|
|
stack.callback(callback_proxy.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
assert _blocked_row(end_user) == [{"blocked": True}]
|
|
|
|
shared: Final = gateway.request("POST", "/customer/unblock", {"user_ids": [end_user]})
|
|
assert shared.status_code == 400, shared.text
|
|
assert shared.json() == {"detail": {"error": "Blocked user check was never set. This call has no effect."}}, (
|
|
shared.text
|
|
)
|
|
assert _blocked_row(end_user) == [{"blocked": True}]
|
|
|
|
callback_response: Final = callback_proxy.request("POST", "/customer/unblock", {"user_ids": [end_user]})
|
|
assert callback_response.status_code == 400, callback_response.text
|
|
assert callback_response.json() == {
|
|
"detail": {"error": "Blocked user check was never set. This call has no effect."}
|
|
}, callback_response.text
|
|
assert _blocked_row(end_user) == [{"blocked": True}]
|
|
|
|
|
|
def test_block_route_returns_the_blocked_records(gateway: Gateway) -> None:
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
with gateway.scenario() as scenario:
|
|
created: Final = gateway.request("POST", "/customer/new", {"user_id": end_user})
|
|
assert created.status_code == 200, created.text
|
|
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
blocked: Final = gateway.request("POST", "/customer/block", {"user_ids": [end_user]})
|
|
assert blocked.status_code == 200, blocked.text
|
|
payload: Final = object_value(blocked.json())
|
|
blocked_users: Final = payload["blocked_users"]
|
|
assert isinstance(blocked_users, list) and len(blocked_users) == 1
|
|
assert object_value(blocked_users[0])["user_id"] == end_user
|
|
assert object_value(blocked_users[0])["blocked"] is True
|
|
|
|
|
|
def test_blocked_end_user_is_refused_without_the_callback(gateway: Gateway) -> None:
|
|
pytest.skip("BUG: blocked end-user request without callback returns 200 instead of 400")
|
|
with gateway.scenario() as scenario:
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
created: Final = gateway.request("POST", "/customer/new", {"user_id": end_user, "blocked": True})
|
|
assert created.status_code == 200, created.text
|
|
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
key: Final = scenario.key(models=[model])
|
|
refused: Final = _chat(gateway, model, key, end_user)
|
|
_assert_blocked_refusal(refused, end_user)
|
|
|
|
|
|
def test_blocked_end_user_header_spelling_is_refused_with_the_callback(gateway: Gateway, tmp_path: Path) -> None:
|
|
pytest.skip("BUG: x-litellm-end-user-id request for a blocked customer returns 200 instead of 400")
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
config: Final = _callback_config(tmp_path, blocked_user_list=[end_user])
|
|
with (
|
|
ExitStack() as stack,
|
|
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
|
):
|
|
proxy: Final = stack.enter_context(owned_proxy(gateway, tmp_path / "listed", {}, config=config))
|
|
with proxy.scenario() as scenario:
|
|
created: Final = proxy.request("POST", "/customer/new", {"user_id": end_user, "max_budget": 100.0})
|
|
assert created.status_code == 200, created.text
|
|
scenario.cleanups.callback(proxy.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
key: Final = scenario.key(models=[model])
|
|
refused: Final = proxy.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": f"blocked {uuid.uuid4().hex}"}]},
|
|
key=key,
|
|
headers={"x-litellm-end-user-id": end_user},
|
|
)
|
|
_assert_blocked_refusal(refused, end_user)
|
|
assert _observed(upstream) == []
|