From 1801c1cf15c924fef3f48d1dd25018457bbcfa01 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 6 May 2026 15:10:21 -0700 Subject: [PATCH] test(agent_session_endpoints): add in-memory Prisma fake + multi-tenant TestClient fixtures --- .../proxy/agent_session_endpoints/conftest.py | 295 ++++++++++++++++++ 1 file changed, 295 insertions(+) create mode 100644 tests/test_litellm/proxy/agent_session_endpoints/conftest.py diff --git a/tests/test_litellm/proxy/agent_session_endpoints/conftest.py b/tests/test_litellm/proxy/agent_session_endpoints/conftest.py new file mode 100644 index 00000000000..3ff63facb82 --- /dev/null +++ b/tests/test_litellm/proxy/agent_session_endpoints/conftest.py @@ -0,0 +1,295 @@ +""" +Shared fixtures for `litellm/proxy/agent_session_endpoints/` tests. + +Provides: +* ``fake_prisma_client`` — an in-memory stand-in for the proxy's Prisma + client. It implements only the methods our endpoints actually call — + no network, no schema. Tests assert against the data structures + directly. +* ``client`` — FastAPI TestClient with all four routers + mounted and ``user_api_key_auth`` overridden to a fixed proxy admin. +* ``other_tenant_client`` — TestClient where the auth dep returns a + different (non-admin) caller, used for cross-tenant isolation tests. +""" + +import os + +import pytest + +# Set a JWT secret BEFORE any module under test is imported. +os.environ.setdefault("LITELLM_AGENT_JWT_SECRET", "test-agent-jwt-secret") +os.environ.setdefault("LITELLM_MASTER_KEY", "sk-1234") + +from collections import defaultdict +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +# --------------------------------------------------------------------------- +# In-memory Prisma stand-in +# --------------------------------------------------------------------------- + + +class _Row: + """Plain object mimicking Prisma row attribute access. + + Missing attributes resolve to ``None`` to mirror Prisma's behavior + of returning a row with optional columns left null. + """ + + def __init__(self, **fields: Any) -> None: + for k, v in fields.items(): + setattr(self, k, v) + + def __getattr__(self, name: str) -> Any: + # Only fires for attribute access that misses the instance dict; + # underscore-prefixed names (e.g. dunders) should error normally. + if name.startswith("_"): + raise AttributeError(name) + return None + + def to_dict(self) -> Dict[str, Any]: + return self.__dict__.copy() + + +def _matches(row: _Row, where: Optional[Dict[str, Any]]) -> bool: + if not where: + return True + for key, expected in where.items(): + actual = getattr(row, key, None) + if isinstance(expected, dict): + if "in" in expected: + if actual not in expected["in"]: + return False + elif "notIn" in expected: + if actual in expected["notIn"]: + return False + elif "lt" in expected: + if actual is None or not (actual < expected["lt"]): + return False + elif "gt" in expected: + if actual is None or not (actual > expected["gt"]): + return False + else: + # Unknown operator dict — fall back to equality on raw dict. + if actual != expected: + return False + else: + if actual != expected: + return False + return True + + +def _order_rows(rows: List[_Row], order: Any) -> List[_Row]: + if not order: + return rows + if isinstance(order, dict): + order = [order] + for o in reversed(order): + for k, direction in o.items(): + rows.sort( + key=lambda r: getattr(r, k) or 0, + reverse=(direction == "desc"), + ) + return rows + + +class _Table: + """In-memory table with the few async methods our endpoints use.""" + + def __init__(self) -> None: + self.rows: List[_Row] = [] + + async def create(self, data: Dict[str, Any]) -> _Row: + # Defaults that real Prisma would apply + now = datetime.now(timezone.utc) + defaults = {"created_at": now, "updated_at": now} + merged = {**defaults, **data} + row = _Row(**merged) + self.rows.append(row) + return row + + async def find_unique(self, where: Dict[str, Any]) -> Optional[_Row]: + for row in self.rows: + if _matches(row, where): + return row + return None + + async def find_first( + self, + where: Optional[Dict[str, Any]] = None, + order: Any = None, + ) -> Optional[_Row]: + results = [r for r in self.rows if _matches(r, where)] + results = _order_rows(results, order) + return results[0] if results else None + + async def find_many( + self, + where: Optional[Dict[str, Any]] = None, + order: Any = None, + take: Optional[int] = None, + skip: Optional[int] = None, + ) -> List[_Row]: + results = [r for r in self.rows if _matches(r, where)] + results = _order_rows(results, order) + if skip: + results = results[skip:] + if take: + results = results[:take] + return results + + async def update(self, where: Dict[str, Any], data: Dict[str, Any]) -> _Row: + for row in self.rows: + if _matches(row, where): + for k, v in data.items(): + setattr(row, k, v) + return row + raise RuntimeError(f"No row to update for {where}") + + async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]): + count = 0 + for row in self.rows: + if _matches(row, where): + for k, v in data.items(): + setattr(row, k, v) + count += 1 + # Mimic Prisma's BatchPayload-ish object. + return _Row(count=count) + + async def delete(self, where: Dict[str, Any]) -> _Row: + for i, row in enumerate(self.rows): + if _matches(row, where): + return self.rows.pop(i) + raise RuntimeError(f"No row to delete for {where}") + + +class FakePrismaClient: + """Drop-in for ``prisma_client`` in ``litellm.proxy.proxy_server``.""" + + def __init__(self) -> None: + self.db = _DB() + + +class _DB: + def __init__(self) -> None: + self.litellm_agent = _Table() + self.litellm_agentsession = _Table() + self.litellm_agentrun = _AgentRunTable() + self.litellm_agentrunevent = _AgentRunEventTable() + + +class _AgentRunTable(_Table): + """Subclass that enforces the (session_id, idempotency_key) unique constraint.""" + + async def create(self, data: Dict[str, Any]) -> _Row: + sid = data.get("session_id") + idem = data.get("idempotency_key") + if idem is not None: + for row in self.rows: + if ( + getattr(row, "session_id", None) == sid + and getattr(row, "idempotency_key", None) == idem + ): + raise RuntimeError("idempotency_collision") + return await super().create(data) + + +class _AgentRunEventTable(_Table): + """Enforces the (run_id, seq) unique constraint.""" + + async def create(self, data: Dict[str, Any]) -> _Row: + rid = data.get("run_id") + seq = data.get("seq") + for row in self.rows: + if getattr(row, "run_id", None) == rid and getattr(row, "seq", None) == seq: + raise RuntimeError("event_seq_collision") + return await super().create(data) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def fake_prisma_client(monkeypatch): + """Patch ``litellm.proxy.proxy_server.prisma_client`` for the duration of + the test with our in-memory stand-in. + """ + from litellm.proxy import proxy_server + + fake = FakePrismaClient() + monkeypatch.setattr(proxy_server, "prisma_client", fake) + return fake + + +def _build_test_app( + role: LitellmUserRoles, api_key: str = "sk-test-caller" +) -> TestClient: + from litellm.proxy.agent_session_endpoints import ( + agent_router, + internal_router, + run_router, + session_router, + ) + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = FastAPI() + app.include_router(agent_router) + app.include_router(session_router) + app.include_router(run_router) + app.include_router(internal_router) + + def _fake_auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_id="test-user", + user_role=role, + api_key=api_key, + team_id=None, + ) + + app.dependency_overrides[user_api_key_auth] = _fake_auth + return TestClient(app) + + +@pytest.fixture +def client(fake_prisma_client): + """TestClient where caller is a non-admin (regular tenant).""" + return _build_test_app(LitellmUserRoles.INTERNAL_USER, api_key="sk-tenant-A") + + +@pytest.fixture +def admin_client(fake_prisma_client): + """TestClient where caller is a proxy admin (sees everything).""" + return _build_test_app(LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin-key") + + +@pytest.fixture +def other_tenant_client(fake_prisma_client): + """TestClient where caller is a different tenant. Used for + cross-tenant isolation tests.""" + return _build_test_app(LitellmUserRoles.INTERNAL_USER, api_key="sk-tenant-B") + + +@pytest.fixture +def noop_provider(monkeypatch): + """Reset the VM provider registry to a fresh ``NoopVMProvider``.""" + from litellm.proxy.agent_session_endpoints.vm_providers import ( + NoopVMProvider, + register_vm_provider, + ) + from litellm.proxy.agent_session_endpoints.vm_providers.registry import ( + reset_vm_provider_registry, + ) + + reset_vm_provider_registry() + provider = NoopVMProvider() + register_vm_provider(provider) + yield provider + reset_vm_provider_registry()