diff --git a/tests/pass_through_unit_tests/test_claude_code_marketplace.py b/tests/pass_through_unit_tests/test_claude_code_marketplace.py index b4ba30e9c7..21387439d8 100644 --- a/tests/pass_through_unit_tests/test_claude_code_marketplace.py +++ b/tests/pass_through_unit_tests/test_claude_code_marketplace.py @@ -6,9 +6,12 @@ Tests: 2. Get marketplace.json (list enabled plugins) """ +import json import os import sys import time +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock import pytest @@ -17,7 +20,6 @@ sys.path.insert(0, os.path.abspath("../..")) import litellm from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.proxy_server import LitellmUserRoles -from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.caching.caching import DualCache from litellm.types.proxy.claude_code_endpoints import RegisterPluginRequest @@ -27,33 +29,118 @@ from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketp get_marketplace, ) -proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + +class MockPluginRecord: + """Mock plugin record that mimics Prisma model behavior.""" + + def __init__(self, name, version, description, manifest_json, enabled=True, created_by=None): + self.id = f"plugin-{name}-{int(time.time())}" + self.name = name + self.version = version + self.description = description + self.manifest_json = manifest_json + self.files_json = "{}" + self.enabled = enabled + self.created_at = datetime.now(timezone.utc) + self.updated_at = datetime.now(timezone.utc) + self.created_by = created_by @pytest.fixture -def prisma_client(): - from litellm.proxy.proxy_cli import append_query_params +def mock_prisma_client(): + """Create a mock PrismaClient that doesn't require Prisma binaries.""" + # In-memory storage for plugins + plugins_store = {} - params = {"connection_limit": 100, "pool_timeout": 60} - database_url = os.getenv("DATABASE_URL") - modified_url = append_query_params(database_url, params) - os.environ["DATABASE_URL"] = modified_url + # Create mock client + mock_client = MagicMock() + mock_client.proxy_logging_obj = MagicMock() - prisma_client = PrismaClient( - database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj - ) + # Mock the db attribute + mock_client.db = MagicMock() - litellm.proxy.proxy_server.litellm_proxy_budget_name = ( - f"litellm-proxy-budget-{time.time()}" - ) + # Mock the plugin table with async methods + mock_table = MagicMock() - return prisma_client + async def find_unique(where): + """Mock find_unique - returns plugin if exists, None otherwise.""" + plugin_name = where.get("name") + return plugins_store.get(plugin_name) + + async def find_many(where=None): + """Mock find_many - returns list of plugins matching where clause.""" + if where is None or where == {}: + return list(plugins_store.values()) + enabled = where.get("enabled") + if enabled is not None: + return [p for p in plugins_store.values() if p.enabled == enabled] + return list(plugins_store.values()) + + async def create(data): + """Mock create - creates a new plugin.""" + plugin_name = data["name"] + manifest = data.get("manifest_json", "{}") + plugin = MockPluginRecord( + name=plugin_name, + version=data.get("version"), + description=data.get("description"), + manifest_json=manifest, + enabled=data.get("enabled", True), + created_by=data.get("created_by"), + ) + plugins_store[plugin_name] = plugin + return plugin + + async def update(where, data): + """Mock update - updates an existing plugin.""" + plugin_name = where.get("name") + if plugin_name not in plugins_store: + raise ValueError(f"Plugin {plugin_name} not found") + plugin = plugins_store[plugin_name] + # Update fields + if "version" in data: + plugin.version = data["version"] + if "description" in data: + plugin.description = data["description"] + if "manifest_json" in data: + plugin.manifest_json = data["manifest_json"] + if "enabled" in data: + plugin.enabled = data["enabled"] + if "updated_at" in data: + plugin.updated_at = data["updated_at"] + return plugin + + async def delete(where): + """Mock delete - deletes a plugin.""" + plugin_name = where.get("name") + if plugin_name in plugins_store: + del plugins_store[plugin_name] + return None + + async def connect(): + """Mock connect - no-op.""" + pass + + # Set up async mocks + mock_table.find_unique = AsyncMock(side_effect=find_unique) + mock_table.find_many = AsyncMock(side_effect=find_many) + mock_table.create = AsyncMock(side_effect=create) + mock_table.update = AsyncMock(side_effect=update) + mock_table.delete = AsyncMock(side_effect=delete) + + mock_client.db.litellm_claudecodeplugintable = mock_table + mock_client.connect = AsyncMock(side_effect=connect) + + # Store plugins_store on the mock for cleanup if needed + mock_client._plugins_store = plugins_store + + return mock_client @pytest.mark.asyncio -async def test_register_plugin(prisma_client): +async def test_register_plugin(mock_prisma_client): """Test registering a plugin in the marketplace.""" - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") await litellm.proxy.proxy_server.prisma_client.connect() @@ -85,16 +172,23 @@ async def test_register_plugin(prisma_client): assert response["plugin"]["version"] == "1.0.0" assert response["plugin"]["enabled"] is True + # Verify the plugin was stored in the mock + stored_plugin = await mock_prisma_client.db.litellm_claudecodeplugintable.find_unique( + where={"name": plugin_name} + ) + assert stored_plugin is not None + assert stored_plugin.name == plugin_name + # Cleanup - delete the plugin - await prisma_client.db.litellm_claudecodeplugintable.delete( + await mock_prisma_client.db.litellm_claudecodeplugintable.delete( where={"name": plugin_name} ) @pytest.mark.asyncio -async def test_get_marketplace(prisma_client): +async def test_get_marketplace(mock_prisma_client): """Test getting marketplace.json with registered plugins.""" - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") await litellm.proxy.proxy_server.prisma_client.connect() @@ -124,7 +218,6 @@ async def test_get_marketplace(prisma_client): response = await get_marketplace() # Response is a JSONResponse, get the body - import json body = json.loads(response.body.decode()) assert body["name"] == "litellm" @@ -140,6 +233,6 @@ async def test_get_marketplace(prisma_client): assert our_plugin["version"] == "2.0.0" # Cleanup - await prisma_client.db.litellm_claudecodeplugintable.delete( + await mock_prisma_client.db.litellm_claudecodeplugintable.delete( where={"name": plugin_name} )