204 lines
8.7 KiB
Python
204 lines
8.7 KiB
Python
|
"""Tests for MCP Forge Server."""
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
from unittest.mock import Mock, AsyncMock, patch
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from mcp_forge.server.server import ForgeServer
|
||
|
|
from mcp_forge.config.schema import ForgeConfig
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_config(tmp_path):
|
||
|
|
"""Mock forge configuration."""
|
||
|
|
config = Mock(spec=ForgeConfig)
|
||
|
|
|
||
|
|
# Server config
|
||
|
|
config.server = Mock()
|
||
|
|
config.server.host = "localhost"
|
||
|
|
config.server.port = 3000
|
||
|
|
config.server.podman_socket = Path("/run/user/1000/podman/podman.sock")
|
||
|
|
|
||
|
|
# Security config
|
||
|
|
config.security = Mock()
|
||
|
|
config.security.audit_log = tmp_path / "audit.log"
|
||
|
|
config.security.max_memory = "2g"
|
||
|
|
config.security.max_timeout = 1800
|
||
|
|
|
||
|
|
# Execution config
|
||
|
|
config.execution = Mock()
|
||
|
|
config.execution.default_backend = "simple"
|
||
|
|
config.execution.default_timeout = 300
|
||
|
|
config.execution.max_timeout = 1800
|
||
|
|
config.execution.default_memory = "512m"
|
||
|
|
config.execution.max_memory = "2g"
|
||
|
|
|
||
|
|
# Sessions config
|
||
|
|
config.sessions = Mock()
|
||
|
|
config.sessions.max_concurrent = 10
|
||
|
|
config.sessions.idle_timeout = 3600
|
||
|
|
|
||
|
|
# Environment builder config
|
||
|
|
config.environment_builder = Mock()
|
||
|
|
config.environment_builder.uv_cache_path = tmp_path / "cache"
|
||
|
|
config.environment_builder.max_build_time = 600
|
||
|
|
config.environment_builder.package_validation = Mock()
|
||
|
|
|
||
|
|
return config
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_server_initialization(mock_config):
|
||
|
|
"""Test that server initializes all components."""
|
||
|
|
with patch('mcp_forge.server.server.AuditLogger'), \
|
||
|
|
patch('mcp_forge.server.server.OperationValidator'), \
|
||
|
|
patch('mcp_forge.server.server.PodmanClient'), \
|
||
|
|
patch('mcp_forge.server.server.SecureContainerManager'), \
|
||
|
|
patch('mcp_forge.server.server.MCPClientManager'), \
|
||
|
|
patch('mcp_forge.server.server.ToolBridgeServer'), \
|
||
|
|
patch('mcp_forge.server.server.ToolInjectionGenerator'), \
|
||
|
|
patch('mcp_forge.server.server.SimpleBackend'), \
|
||
|
|
patch('mcp_forge.server.server.JupyterKernelManager'), \
|
||
|
|
patch('mcp_forge.server.server.JupyterBackend'), \
|
||
|
|
patch('mcp_forge.server.server.SessionManager'), \
|
||
|
|
patch('mcp_forge.server.server.EnvironmentBuilder'):
|
||
|
|
|
||
|
|
server = ForgeServer(config=mock_config)
|
||
|
|
|
||
|
|
# Verify server was created
|
||
|
|
assert server is not None
|
||
|
|
assert server.config == mock_config
|
||
|
|
|
||
|
|
# Verify components were initialized
|
||
|
|
assert hasattr(server, 'audit_logger')
|
||
|
|
assert hasattr(server, 'operation_validator')
|
||
|
|
assert hasattr(server, 'podman_client')
|
||
|
|
assert hasattr(server, 'container_manager')
|
||
|
|
assert hasattr(server, 'client_manager')
|
||
|
|
assert hasattr(server, 'bridge_server')
|
||
|
|
assert hasattr(server, 'simple_backend')
|
||
|
|
assert hasattr(server, 'jupyter_backend')
|
||
|
|
assert hasattr(server, 'environment_builder')
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_tools_registration(mock_config):
|
||
|
|
"""Test that tools are registered with the server."""
|
||
|
|
with patch('mcp_forge.server.server.AuditLogger'), \
|
||
|
|
patch('mcp_forge.server.server.OperationValidator'), \
|
||
|
|
patch('mcp_forge.server.server.PodmanClient'), \
|
||
|
|
patch('mcp_forge.server.server.SecureContainerManager'), \
|
||
|
|
patch('mcp_forge.server.server.MCPClientManager'), \
|
||
|
|
patch('mcp_forge.server.server.ToolBridgeServer'), \
|
||
|
|
patch('mcp_forge.server.server.ToolInjectionGenerator'), \
|
||
|
|
patch('mcp_forge.server.server.SimpleBackend'), \
|
||
|
|
patch('mcp_forge.server.server.JupyterKernelManager'), \
|
||
|
|
patch('mcp_forge.server.server.JupyterBackend'), \
|
||
|
|
patch('mcp_forge.server.server.SessionManager'), \
|
||
|
|
patch('mcp_forge.server.server.EnvironmentBuilder'):
|
||
|
|
|
||
|
|
server = ForgeServer(config=mock_config)
|
||
|
|
|
||
|
|
# Verify tool instances were created
|
||
|
|
assert hasattr(server, 'execute_python_tool')
|
||
|
|
assert hasattr(server, 'document_state_tool')
|
||
|
|
assert hasattr(server, 'build_environment_tool')
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_resources_registration(mock_config):
|
||
|
|
"""Test that resources are registered with the server."""
|
||
|
|
with patch('mcp_forge.server.server.AuditLogger'), \
|
||
|
|
patch('mcp_forge.server.server.OperationValidator'), \
|
||
|
|
patch('mcp_forge.server.server.PodmanClient'), \
|
||
|
|
patch('mcp_forge.server.server.SecureContainerManager'), \
|
||
|
|
patch('mcp_forge.server.server.MCPClientManager'), \
|
||
|
|
patch('mcp_forge.server.server.ToolBridgeServer'), \
|
||
|
|
patch('mcp_forge.server.server.ToolInjectionGenerator'), \
|
||
|
|
patch('mcp_forge.server.server.SimpleBackend'), \
|
||
|
|
patch('mcp_forge.server.server.JupyterKernelManager'), \
|
||
|
|
patch('mcp_forge.server.server.JupyterBackend'), \
|
||
|
|
patch('mcp_forge.server.server.SessionManager'), \
|
||
|
|
patch('mcp_forge.server.server.EnvironmentBuilder'):
|
||
|
|
|
||
|
|
server = ForgeServer(config=mock_config)
|
||
|
|
|
||
|
|
# Verify resource handler was created
|
||
|
|
assert hasattr(server, 'resource_handler')
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_server_shutdown(mock_config):
|
||
|
|
"""Test that server shuts down gracefully."""
|
||
|
|
with patch('mcp_forge.server.server.AuditLogger') as mock_audit, \
|
||
|
|
patch('mcp_forge.server.server.OperationValidator'), \
|
||
|
|
patch('mcp_forge.server.server.PodmanClient'), \
|
||
|
|
patch('mcp_forge.server.server.SecureContainerManager'), \
|
||
|
|
patch('mcp_forge.server.server.MCPClientManager') as mock_client_mgr, \
|
||
|
|
patch('mcp_forge.server.server.ToolBridgeServer') as mock_bridge, \
|
||
|
|
patch('mcp_forge.server.server.ToolInjectionGenerator'), \
|
||
|
|
patch('mcp_forge.server.server.SimpleBackend'), \
|
||
|
|
patch('mcp_forge.server.server.JupyterKernelManager'), \
|
||
|
|
patch('mcp_forge.server.server.JupyterBackend'), \
|
||
|
|
patch('mcp_forge.server.server.SessionManager'), \
|
||
|
|
patch('mcp_forge.server.server.EnvironmentBuilder'):
|
||
|
|
|
||
|
|
# Setup mocks
|
||
|
|
mock_client_mgr.return_value.shutdown = AsyncMock()
|
||
|
|
mock_bridge_instance = Mock()
|
||
|
|
mock_bridge_instance.stop = Mock() # Not async
|
||
|
|
mock_bridge.return_value = mock_bridge_instance
|
||
|
|
mock_audit_instance = Mock()
|
||
|
|
mock_audit.return_value = mock_audit_instance
|
||
|
|
|
||
|
|
server = ForgeServer(config=mock_config)
|
||
|
|
|
||
|
|
# Shutdown server
|
||
|
|
await server.shutdown()
|
||
|
|
|
||
|
|
# Verify cleanup was called
|
||
|
|
server.client_manager.shutdown.assert_called_once()
|
||
|
|
server.bridge_server.stop.assert_called_once()
|
||
|
|
server.audit_logger.log.assert_called()
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_server_component_initialization_order(mock_config):
|
||
|
|
"""Test that components are initialized in correct order."""
|
||
|
|
init_order = []
|
||
|
|
|
||
|
|
def track_init(name):
|
||
|
|
def decorator(cls):
|
||
|
|
original_init = cls.__init__
|
||
|
|
def new_init(self, *args, **kwargs):
|
||
|
|
init_order.append(name)
|
||
|
|
return original_init(self, *args, **kwargs)
|
||
|
|
cls.__init__ = new_init
|
||
|
|
return cls
|
||
|
|
return decorator
|
||
|
|
|
||
|
|
with patch('mcp_forge.server.server.AuditLogger') as mock_audit, \
|
||
|
|
patch('mcp_forge.server.server.OperationValidator') as mock_validator, \
|
||
|
|
patch('mcp_forge.server.server.PodmanClient') as mock_podman, \
|
||
|
|
patch('mcp_forge.server.server.SecureContainerManager') as mock_container, \
|
||
|
|
patch('mcp_forge.server.server.MCPClientManager'), \
|
||
|
|
patch('mcp_forge.server.server.ToolBridgeServer'), \
|
||
|
|
patch('mcp_forge.server.server.ToolInjectionGenerator'), \
|
||
|
|
patch('mcp_forge.server.server.SimpleBackend'), \
|
||
|
|
patch('mcp_forge.server.server.JupyterKernelManager'), \
|
||
|
|
patch('mcp_forge.server.server.JupyterBackend'), \
|
||
|
|
patch('mcp_forge.server.server.SessionManager'), \
|
||
|
|
patch('mcp_forge.server.server.EnvironmentBuilder'):
|
||
|
|
|
||
|
|
mock_audit.side_effect = lambda *args, **kwargs: init_order.append('audit_logger') or Mock()
|
||
|
|
mock_validator.side_effect = lambda *args, **kwargs: init_order.append('operation_validator') or Mock()
|
||
|
|
mock_podman.side_effect = lambda *args, **kwargs: init_order.append('podman_client') or Mock()
|
||
|
|
mock_container.side_effect = lambda *args, **kwargs: init_order.append('container_manager') or Mock()
|
||
|
|
|
||
|
|
server = ForgeServer(config=mock_config)
|
||
|
|
|
||
|
|
# Verify security components are initialized first
|
||
|
|
assert init_order.index('audit_logger') < init_order.index('podman_client')
|
||
|
|
assert init_order.index('operation_validator') < init_order.index('podman_client')
|
||
|
|
assert init_order.index('podman_client') < init_order.index('container_manager')
|