forked from cardosofelipe/fast-next-template
Felipe Cardoso
520c06175e
Improved code readability and uniformity by standardizing line breaks, indentation, and inline conditions across safety-related services, models, and tests, including content filters, validation rules, and emergency controls.
273 lines
8.9 KiB
Python
273 lines
8.9 KiB
Python
"""
|
|
Tests for MCP Server Registry
|
|
"""
|
|
|
|
import pytest
|
|
|
|
from app.services.mcp.config import MCPConfig, MCPServerConfig
|
|
from app.services.mcp.exceptions import MCPServerNotFoundError
|
|
from app.services.mcp.registry import (
|
|
MCPServerRegistry,
|
|
ServerCapabilities,
|
|
get_registry,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def reset_registry():
|
|
"""Reset the singleton registry before and after each test."""
|
|
MCPServerRegistry.reset_instance()
|
|
yield
|
|
MCPServerRegistry.reset_instance()
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_config():
|
|
"""Create a sample MCP configuration."""
|
|
return MCPConfig(
|
|
mcp_servers={
|
|
"server-1": MCPServerConfig(
|
|
url="http://server1:8000",
|
|
timeout=30,
|
|
enabled=True,
|
|
),
|
|
"server-2": MCPServerConfig(
|
|
url="http://server2:8000",
|
|
timeout=60,
|
|
enabled=True,
|
|
),
|
|
"disabled-server": MCPServerConfig(
|
|
url="http://disabled:8000",
|
|
enabled=False,
|
|
),
|
|
}
|
|
)
|
|
|
|
|
|
class TestServerCapabilities:
|
|
"""Tests for ServerCapabilities class."""
|
|
|
|
def test_empty_capabilities(self):
|
|
"""Test creating empty capabilities."""
|
|
caps = ServerCapabilities()
|
|
assert caps.tools == []
|
|
assert caps.resources == []
|
|
assert caps.prompts == []
|
|
assert caps.is_loaded is False
|
|
assert caps.tool_names == []
|
|
|
|
def test_capabilities_with_tools(self):
|
|
"""Test capabilities with tools."""
|
|
caps = ServerCapabilities(
|
|
tools=[
|
|
{"name": "tool1", "description": "Tool 1"},
|
|
{"name": "tool2", "description": "Tool 2"},
|
|
]
|
|
)
|
|
assert len(caps.tools) == 2
|
|
assert caps.tool_names == ["tool1", "tool2"]
|
|
|
|
def test_mark_loaded(self):
|
|
"""Test marking capabilities as loaded."""
|
|
caps = ServerCapabilities()
|
|
assert caps.is_loaded is False
|
|
assert caps._load_time is None
|
|
|
|
caps.mark_loaded()
|
|
assert caps.is_loaded is True
|
|
assert caps._load_time is not None
|
|
|
|
|
|
class TestMCPServerRegistry:
|
|
"""Tests for MCPServerRegistry singleton."""
|
|
|
|
def test_singleton_pattern(self, reset_registry):
|
|
"""Test that registry is a singleton."""
|
|
registry1 = MCPServerRegistry()
|
|
registry2 = MCPServerRegistry()
|
|
assert registry1 is registry2
|
|
|
|
def test_get_instance(self, reset_registry):
|
|
"""Test get_instance class method."""
|
|
registry = MCPServerRegistry.get_instance()
|
|
assert registry is MCPServerRegistry()
|
|
|
|
def test_reset_instance(self, reset_registry):
|
|
"""Test resetting singleton instance."""
|
|
registry1 = MCPServerRegistry()
|
|
MCPServerRegistry.reset_instance()
|
|
registry2 = MCPServerRegistry()
|
|
assert registry1 is not registry2
|
|
|
|
def test_load_config(self, reset_registry, sample_config):
|
|
"""Test loading configuration."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
assert len(registry.list_servers()) == 3
|
|
assert "server-1" in registry.list_servers()
|
|
assert "server-2" in registry.list_servers()
|
|
assert "disabled-server" in registry.list_servers()
|
|
|
|
def test_list_enabled_servers(self, reset_registry, sample_config):
|
|
"""Test listing only enabled servers."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
enabled = registry.list_enabled_servers()
|
|
assert len(enabled) == 2
|
|
assert "server-1" in enabled
|
|
assert "server-2" in enabled
|
|
assert "disabled-server" not in enabled
|
|
|
|
def test_register(self, reset_registry):
|
|
"""Test registering a new server."""
|
|
registry = MCPServerRegistry()
|
|
config = MCPServerConfig(url="http://new:8000")
|
|
|
|
registry.register("new-server", config)
|
|
assert "new-server" in registry.list_servers()
|
|
assert registry.get("new-server").url == "http://new:8000"
|
|
|
|
def test_unregister(self, reset_registry, sample_config):
|
|
"""Test unregistering a server."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
assert registry.unregister("server-1") is True
|
|
assert "server-1" not in registry.list_servers()
|
|
|
|
# Unregistering non-existent server returns False
|
|
assert registry.unregister("nonexistent") is False
|
|
|
|
def test_get(self, reset_registry, sample_config):
|
|
"""Test getting server config."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
config = registry.get("server-1")
|
|
assert config.url == "http://server1:8000"
|
|
assert config.timeout == 30
|
|
|
|
def test_get_not_found(self, reset_registry, sample_config):
|
|
"""Test getting non-existent server raises error."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
with pytest.raises(MCPServerNotFoundError) as exc_info:
|
|
registry.get("nonexistent")
|
|
|
|
assert exc_info.value.server_name == "nonexistent"
|
|
assert "server-1" in exc_info.value.available_servers
|
|
|
|
def test_get_or_none(self, reset_registry, sample_config):
|
|
"""Test get_or_none method."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
config = registry.get_or_none("server-1")
|
|
assert config is not None
|
|
|
|
config = registry.get_or_none("nonexistent")
|
|
assert config is None
|
|
|
|
def test_get_all_configs(self, reset_registry, sample_config):
|
|
"""Test getting all configs."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
configs = registry.get_all_configs()
|
|
assert len(configs) == 3
|
|
|
|
def test_get_enabled_configs(self, reset_registry, sample_config):
|
|
"""Test getting enabled configs."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
configs = registry.get_enabled_configs()
|
|
assert len(configs) == 2
|
|
assert "disabled-server" not in configs
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_capabilities(self, reset_registry, sample_config):
|
|
"""Test getting server capabilities."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
# Initially empty capabilities
|
|
caps = await registry.get_capabilities("server-1")
|
|
assert caps.is_loaded is False
|
|
|
|
def test_set_capabilities(self, reset_registry, sample_config):
|
|
"""Test setting server capabilities."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
registry.set_capabilities(
|
|
"server-1",
|
|
tools=[{"name": "tool1"}, {"name": "tool2"}],
|
|
resources=[{"name": "resource1"}],
|
|
)
|
|
|
|
caps = registry._capabilities["server-1"]
|
|
assert len(caps.tools) == 2
|
|
assert len(caps.resources) == 1
|
|
assert caps.is_loaded is True
|
|
|
|
def test_find_server_for_tool(self, reset_registry, sample_config):
|
|
"""Test finding server that provides a tool."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
registry.set_capabilities(
|
|
"server-1",
|
|
tools=[{"name": "tool1"}, {"name": "tool2"}],
|
|
)
|
|
registry.set_capabilities(
|
|
"server-2",
|
|
tools=[{"name": "tool3"}],
|
|
)
|
|
|
|
assert registry.find_server_for_tool("tool1") == "server-1"
|
|
assert registry.find_server_for_tool("tool3") == "server-2"
|
|
assert registry.find_server_for_tool("unknown") is None
|
|
|
|
def test_get_all_tools(self, reset_registry, sample_config):
|
|
"""Test getting all tools from all servers."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
registry.set_capabilities(
|
|
"server-1",
|
|
tools=[{"name": "tool1"}],
|
|
)
|
|
registry.set_capabilities(
|
|
"server-2",
|
|
tools=[{"name": "tool2"}, {"name": "tool3"}],
|
|
)
|
|
|
|
all_tools = registry.get_all_tools()
|
|
assert len(all_tools) == 2
|
|
assert len(all_tools["server-1"]) == 1
|
|
assert len(all_tools["server-2"]) == 2
|
|
|
|
def test_global_config_property(self, reset_registry, sample_config):
|
|
"""Test accessing global config."""
|
|
registry = MCPServerRegistry()
|
|
registry.load_config(sample_config)
|
|
|
|
global_config = registry.global_config
|
|
assert global_config is not None
|
|
assert len(global_config.mcp_servers) == 3
|
|
|
|
|
|
class TestGetRegistry:
|
|
"""Tests for get_registry convenience function."""
|
|
|
|
def test_get_registry_returns_singleton(self, reset_registry):
|
|
"""Test that get_registry returns the singleton."""
|
|
registry1 = get_registry()
|
|
registry2 = get_registry()
|
|
assert registry1 is registry2
|
|
assert registry1 is MCPServerRegistry()
|