mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
425c8d37fc
Forwarding limit to OpenAI made the ownership filter cut the page down after the fact, so a key that owned an older container got an empty first page and its cursor never moved. Non-admin lists now walk upstream pages of 100 until they have enough owned containers (or five pages), trim to the requested limit, and report first_id, last_id and has_more off what the caller keeps. Also assigns tests/test_litellm/proxy/container_endpoints to a CI shard.
134 lines
4.6 KiB
Python
134 lines
4.6 KiB
Python
import sys
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|
from litellm.proxy.container_endpoints import endpoints, ownership
|
|
from litellm.types.containers.main import ContainerListResponse, ContainerObject
|
|
|
|
PROXY_SERVER_STUB = SimpleNamespace(
|
|
general_settings={},
|
|
prisma_client=None,
|
|
llm_router=None,
|
|
proxy_config=None,
|
|
proxy_logging_obj=None,
|
|
select_data_generator=None,
|
|
user_api_base=None,
|
|
user_max_tokens=None,
|
|
user_model=None,
|
|
user_request_timeout=None,
|
|
user_temperature=None,
|
|
version="test",
|
|
)
|
|
ADMIN = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN)
|
|
NON_ADMIN = UserAPIKeyAuth(user_id="user-1")
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clear_allowed_container_ids_cache():
|
|
ownership._ALLOWED_CONTAINER_IDS_CACHE.cache_dict.clear()
|
|
ownership._ALLOWED_CONTAINER_IDS_CACHE.ttl_dict.clear()
|
|
yield
|
|
ownership._ALLOWED_CONTAINER_IDS_CACHE.cache_dict.clear()
|
|
ownership._ALLOWED_CONTAINER_IDS_CACHE.ttl_dict.clear()
|
|
|
|
|
|
def _client(auth: UserAPIKeyAuth) -> TestClient:
|
|
app = FastAPI()
|
|
app.include_router(endpoints.router)
|
|
app.dependency_overrides[user_api_key_auth] = lambda: auth
|
|
return TestClient(app)
|
|
|
|
|
|
def _container(container_id: str) -> ContainerObject:
|
|
return ContainerObject(id=container_id, object="container", created_at=1, status="active")
|
|
|
|
|
|
def _page(*container_ids: str, has_more: bool) -> ContainerListResponse:
|
|
return ContainerListResponse(
|
|
object="list",
|
|
data=[_container(container_id) for container_id in container_ids],
|
|
has_more=has_more,
|
|
)
|
|
|
|
|
|
def _upstream_pages(monkeypatch, pages_by_after) -> MagicMock:
|
|
processor_cls = MagicMock(
|
|
side_effect=lambda data: SimpleNamespace(
|
|
base_process_llm_request=AsyncMock(return_value=pages_by_after[data["after"]])
|
|
)
|
|
)
|
|
monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", processor_cls)
|
|
return processor_cls
|
|
|
|
|
|
def _forwarded_pages(processor_cls: MagicMock):
|
|
return [(call.kwargs["data"]["after"], call.kwargs["data"]["limit"]) for call in processor_cls.call_args_list]
|
|
|
|
|
|
def test_list_containers_forwards_typed_pagination_params_for_admins(monkeypatch):
|
|
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", PROXY_SERVER_STUB)
|
|
processor_cls = _upstream_pages(monkeypatch, {"cntr_prev": _page("cntr_next", has_more=True)})
|
|
|
|
response = _client(ADMIN).get(
|
|
"/v1/containers",
|
|
params={"limit": "1", "order": "desc", "after": "cntr_prev"},
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert [item["id"] for item in response.json()["data"]] == ["cntr_next"]
|
|
assert response.json()["has_more"] is True
|
|
assert _forwarded_pages(processor_cls) == [("cntr_prev", 1)]
|
|
assert processor_cls.call_args.kwargs["data"]["order"] == "desc"
|
|
|
|
|
|
def test_list_containers_rejects_a_non_integer_limit(monkeypatch):
|
|
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", PROXY_SERVER_STUB)
|
|
processor_cls = _upstream_pages(monkeypatch, {})
|
|
|
|
response = _client(ADMIN).get(
|
|
"/v1/containers",
|
|
params={"limit": "abc"},
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
|
|
assert response.status_code == 422
|
|
processor_cls.assert_not_called()
|
|
|
|
|
|
def test_list_containers_pages_upstream_until_non_admin_keys_see_their_containers(monkeypatch):
|
|
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", PROXY_SERVER_STUB)
|
|
table = AsyncMock()
|
|
table.find_many.return_value = [SimpleNamespace(model_object_id="container:openai:cntr_owned")]
|
|
monkeypatch.setattr(
|
|
ownership,
|
|
"_get_prisma_client",
|
|
AsyncMock(return_value=SimpleNamespace(db=SimpleNamespace(litellm_managedobjecttable=table))),
|
|
)
|
|
processor_cls = _upstream_pages(
|
|
monkeypatch,
|
|
{
|
|
None: _page("cntr_other", has_more=True),
|
|
"cntr_other": _page("cntr_owned", has_more=False),
|
|
},
|
|
)
|
|
|
|
response = _client(NON_ADMIN).get(
|
|
"/v1/containers",
|
|
params={"limit": "1"},
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
body = response.json()
|
|
assert [item["id"] for item in body["data"]] == ["cntr_owned"]
|
|
assert body["first_id"] == "cntr_owned"
|
|
assert body["last_id"] == "cntr_owned"
|
|
assert body["has_more"] is False
|
|
assert _forwarded_pages(processor_cls) == [(None, 100), ("cntr_other", 100)]
|