From af4c6de05ee80ebf07a557f835ac94a187249e6a Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 7 May 2026 10:06:56 -0700 Subject: [PATCH] test(v2 managed agents): unit tests for POST /v2/agents --- .../managed_agents/test_agents.py | 313 ++++++++++++++++++ 1 file changed, 313 insertions(+) create mode 100644 tests/test_litellm/managed_agents/test_agents.py diff --git a/tests/test_litellm/managed_agents/test_agents.py b/tests/test_litellm/managed_agents/test_agents.py new file mode 100644 index 00000000000..53cd064c922 --- /dev/null +++ b/tests/test_litellm/managed_agents/test_agents.py @@ -0,0 +1,313 @@ +"""Unit tests for ``POST /v2/agents`` (LIT-2922). + +Auth pattern: FastAPI ``app.dependency_overrides`` for ``user_api_key_auth``. + +The router is included onto the proxy ``app`` inside a fixture so this test +file does not require Wave 3's wiring in ``proxy_server.py`` to be present +yet. The fixture is idempotent — re-including a router on the same app is +a no-op for FastAPI's matcher (it just registers another mount); we de-dup +by checking ``app.router.routes`` first. +""" + +import os +import sys +import types +from datetime import datetime +from typing import Any, Dict, Optional +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi.testclient import TestClient + +import litellm.proxy.proxy_server as ps +from litellm.managed_agents.endpoints.agents import router as agents_router +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.proxy_server import app + +sys.path.insert(0, os.path.abspath("../../../")) + + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- + + +def _ensure_router_mounted() -> None: + """Mount the agents router onto ``app`` exactly once. + + Wave 3 owns the real registration in ``proxy_server.py``. For tests we + mount it here so we can hit ``POST /v2/agents`` via ``TestClient``. + """ + paths = {getattr(r, "path", None) for r in app.router.routes} + if "/v2/agents" not in paths: + app.include_router(agents_router) + + +class _FakeAgentTable: + """In-memory stand-in for ``prisma_client.db.litellm_managedagent``. + + Only the methods called by Wave 1's ``db.py`` helpers are implemented: + ``create`` and ``find_first``. Both are ``AsyncMock`` so call assertions + work as expected. + """ + + def __init__(self) -> None: + self.rows: Dict[str, Dict[str, Any]] = {} + self.create = AsyncMock(side_effect=self._create) + self.find_first = AsyncMock(side_effect=self._find_first) + + async def _create(self, *, data: Dict[str, Any]) -> types.SimpleNamespace: + # Mirror Prisma's behavior: store + return a row-like object. + # ``prisma.Json(...)`` wraps the config dict — unwrap it for storage + # so the test assertions can inspect raw dicts. + config = data.get("config") + if config is not None and hasattr(config, "data"): + stored_config: Any = config.data + elif isinstance(config, dict): + stored_config = config + else: + # ``prisma.Json`` instances expose the original dict through + # ``.data`` in real Prisma, but the stub used in CI may pass + # through. Fall back to the raw value. + stored_config = config + + row = {**data, "config": stored_config} + self.rows[data["id"]] = row + # Return a model_dump-able stand-in. + return types.SimpleNamespace(model_dump=lambda r=row: dict(r)) + + async def _find_first( + self, *, where: Dict[str, Any] + ) -> Optional[types.SimpleNamespace]: + for row in self.rows.values(): + if all(row.get(k) == v for k, v in where.items()): + return types.SimpleNamespace(model_dump=lambda r=row: dict(r)) + return None + + +@pytest.fixture +def client_and_mocks(monkeypatch): + _ensure_router_mounted() + + # 1. Stub prisma client at the module ``ps`` (proxy_server) sees. + fake_table = _FakeAgentTable() + fake_db = types.SimpleNamespace(litellm_managedagent=fake_table) + fake_prisma = MagicMock() + fake_prisma.db = fake_db + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + + # 2. Stub ``prisma.Json`` so ``db.insert_agent`` can wrap the config. + fake_prisma_module = types.ModuleType("prisma") + + class _Json: + def __init__(self, data: Any) -> None: + self.data = data + + fake_prisma_module.Json = _Json # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "prisma", fake_prisma_module) + + # 3. Override the auth dependency with a fixed user. + fake_user = UserAPIKeyAuth( + user_id="user_xyz", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: fake_user + + client = TestClient(app) + + yield client, fake_table, fake_user + + app.dependency_overrides.clear() + + +def _valid_payload(name: str = "code-reviewer") -> Dict[str, Any]: + return { + "name": name, + "config": { + "model": "anthropic/claude-opus-4", + "system_prompt": "You are a senior engineer reviewing code.", + "tools": ["read", "grep", "bash"], + "litellm_api_key": "sk-supersecret123", + "litellm_base_url": "http://localhost:4000", + }, + } + + +# --------------------------------------------------------------------------- +# Happy path +# --------------------------------------------------------------------------- + + +def test_create_agent_success(client_and_mocks): + client, fake_table, fake_user = client_and_mocks + + payload = _valid_payload() + resp = client.post("/v2/agents", json=payload) + + assert resp.status_code == 200, resp.text + body = resp.json() + + assert body["name"] == payload["name"] + assert body["created_by"] == fake_user.user_id + assert body["id"].startswith("agt_") + # Spec §6.1: id is "agt_" + 32-char uuid hex (no dashes) → length 36. + assert len(body["id"]) == 36 + + # Config: passthrough fields preserved, api_key masked. + cfg = body["config"] + assert cfg["model"] == payload["config"]["model"] + assert cfg["system_prompt"] == payload["config"]["system_prompt"] + assert cfg["tools"] == payload["config"]["tools"] + assert cfg["litellm_base_url"] == payload["config"]["litellm_base_url"] + # Mask: first 4 chars + "****" → "sk-s****" + assert cfg["litellm_api_key"] == "sk-s****" + assert payload["config"]["litellm_api_key"] not in resp.text + + # Timestamps are present and ISO-formatted. + assert "created_at" in body and "updated_at" in body + datetime.fromisoformat(body["created_at"].replace("Z", "+00:00")) + datetime.fromisoformat(body["updated_at"].replace("Z", "+00:00")) + + # DB was called exactly once for create. + fake_table.create.assert_awaited_once() + create_kwargs = fake_table.create.await_args.kwargs + inserted = create_kwargs["data"] + # The raw key is what gets persisted (masking is on-read only). + assert ( + inserted["config"].data["litellm_api_key"] + == payload["config"]["litellm_api_key"] + ) + assert inserted["created_by"] == fake_user.user_id + assert inserted["name"] == payload["name"] + + +# --------------------------------------------------------------------------- +# Duplicate names +# --------------------------------------------------------------------------- + + +def test_create_agent_duplicate_name_same_user_returns_409(client_and_mocks): + client, fake_table, _ = client_and_mocks + + payload = _valid_payload(name="code-reviewer") + first = client.post("/v2/agents", json=payload) + assert first.status_code == 200, first.text + + second = client.post("/v2/agents", json=payload) + assert second.status_code == 409, second.text + detail = second.json()["detail"] + assert "code-reviewer" in detail + assert "already exists" in detail.lower() + + # Only the first call should have hit ``create``. + assert fake_table.create.await_count == 1 + + +def test_create_agent_same_name_different_user_returns_200( + client_and_mocks, monkeypatch +): + client, fake_table, _ = client_and_mocks + + # First user creates "code-reviewer". + first_payload = _valid_payload(name="code-reviewer") + first = client.post("/v2/agents", json=first_payload) + assert first.status_code == 200, first.text + + # Swap auth to a different user; same name must succeed. + other_user = UserAPIKeyAuth( + user_id="user_other", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: other_user + + second = client.post("/v2/agents", json=first_payload) + assert second.status_code == 200, second.text + assert second.json()["created_by"] == "user_other" + + # Two distinct rows should now be present. + assert fake_table.create.await_count == 2 + created_bys = {row["created_by"] for row in fake_table.rows.values()} + assert created_bys == {"user_xyz", "user_other"} + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + + +def test_create_agent_missing_name_returns_422(client_and_mocks): + client, fake_table, _ = client_and_mocks + + payload = _valid_payload() + payload.pop("name") + resp = client.post("/v2/agents", json=payload) + + assert resp.status_code == 422, resp.text + detail = resp.json()["detail"] + assert isinstance(detail, list) + assert any("name" in (err.get("loc") or []) for err in detail) + fake_table.create.assert_not_awaited() + + +def test_create_agent_missing_config_model_returns_422(client_and_mocks): + client, fake_table, _ = client_and_mocks + + payload = _valid_payload() + payload["config"].pop("model") + resp = client.post("/v2/agents", json=payload) + + assert resp.status_code == 422, resp.text + detail = resp.json()["detail"] + assert isinstance(detail, list) + assert any( + "model" in (err.get("loc") or []) and "config" in (err.get("loc") or []) + for err in detail + ) + fake_table.create.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Auth +# --------------------------------------------------------------------------- + + +def test_create_agent_no_auth_returns_401(client_and_mocks, monkeypatch): + """Without an ``Authorization`` header and a master key set, real auth rejects. + + The proxy's ``user_api_key_auth`` is permissive when ``master_key is None`` + — it returns a default ``UserAPIKeyAuth`` so local dev doesn't need a key. + To actually exercise the rejection path we set a master key, drop the + dependency override, and send no header. + """ + client, fake_table, _ = client_and_mocks + + # Set a master key so unauth'd requests get rejected by the real chain. + monkeypatch.setattr(ps, "master_key", "sk-test-master") + + # Drop the override so the real user_api_key_auth runs on this request. + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + resp = client.post("/v2/agents", json=_valid_payload(name="no-auth")) + + # The real auth chain rejects unauthenticated calls. Proxy maps the raised + # exception to ProxyException → 401 in most versions, but some 4xx is fine + # — what matters is that it's an auth failure, NOT a successful 2xx, and + # the DB was never touched. + assert resp.status_code in (400, 401, 403), resp.text + fake_table.create.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# DB-not-connected guard +# --------------------------------------------------------------------------- + + +def test_create_agent_db_not_connected_returns_500(client_and_mocks, monkeypatch): + client, _, _ = client_and_mocks + + monkeypatch.setattr(ps, "prisma_client", None) + + resp = client.post("/v2/agents", json=_valid_payload(name="no-db")) + assert resp.status_code == 500, resp.text + detail = resp.json()["detail"] + assert "DB not connected" in detail["error"]