ds_mcp_template/tests/unit/test_server.py

"""
Unit tests for MCP server components
"""
import pytest
import asyncio
import tempfile
import os
from unittest.mock import Mock, AsyncMock, patch, MagicMock
from pathlib import Path
from mcp_template.core.types import MCPTool, MCPResource, MCPPrompt, MCPServerConfig, TransportType
# FastMCPServer not available, using ModularMCPServer for testing
from mcp_template.server.server_factory import MCPServerFactory
from mcp_template.server.modular_server import ModularMCPServer
from mcp_template.server.tools.tool_registry import ServerToolRegistry
from mcp_template.server.prompts.prompt_registry import ServerPromptRegistry
from mcp_template.server.resources.resource_registry import ServerResourceRegistry
from mcp_template.server.tools.base_tool import BaseServerTool, ContextAwareTool
from mcp_template.server.prompts.base_prompt import BaseServerPrompt
from mcp_template.server.resources.base_resource import BaseServerResource


class TestModularMCPServerBasic:
    """Test basic ModularMCP server functionality (replacing FastMCPServer tests)"""

    @pytest.fixture
    def sample_tool(self):
        """Create a sample tool for testing"""
        async def add_handler(a: int, b: int) -> int:
            return a + b

        return MCPTool(
            name="add",
            description="Add two numbers",
            input_schema={
                "type": "object",
                "properties": {
                    "a": {"type": "number"},
                    "b": {"type": "number"}
                },
                "required": ["a", "b"]
            },
            handler=add_handler
        )

    @pytest.fixture
    def server_config(self):
        """Create a sample server config"""
        return MCPServerConfig(
            name="Test Server",
            transport=TransportType.STDIO,
            host="localhost",
            port=8050
        )

    @pytest.mark.asyncio
    async def test_server_initialization(self, server_config, sample_tool):
        """Test server initialization"""
        # Use factory to create server
        server = MCPServerFactory.create_server(
            name="Test Server",
            transport="stdio"
        )

        # Server should not be initialized yet
        assert not server._initialized

        # Initialize server
        await server.initialize()

        # Server should be initialized
        assert server._initialized

    @pytest.mark.asyncio
    async def test_server_info(self, server_config):
        """Test getting server information"""
        server = MCPServerFactory.create_server(
            name="Test Server",
            transport="stdio"
        )

        # Test server info before initialization
        info = server.get_server_info()
        assert info["name"] == "Test Server"
        assert info["host"] == "0.0.0.0"
        assert info["port"] == 8050

    def test_server_run_method(self, server_config):
        """Test server run method"""
        server = MCPServerFactory.create_server(
            name="Test Server",
            transport="stdio"
        )

        # Mock the run method to avoid actually starting the server
        with patch.object(server.mcp, 'run') as mock_run:
            server.run("stdio")
            mock_run.assert_called_once_with(transport="stdio")



class TestMCPServerFactory:
    """Test MCP server factory"""

    def test_create_server_basic(self):
        """Test creating a basic server"""
        server = MCPServerFactory.create_server(
            name="Test Server",
            transport="stdio"
        )

        assert isinstance(server, ModularMCPServer)
        assert server.name == "Test Server"
        assert server.host == "0.0.0.0"
        assert server.port == 8050

    def test_create_calculator_server(self):
        """Test creating a calculator server"""
        server = MCPServerFactory.create_basic_calculator_server(
            name="Calculator Server",
            transport="sse",
            port=8080
        )

        assert isinstance(server, ModularMCPServer)
        assert server.name == "Calculator Server"
        assert server.host == "0.0.0.0"
        assert server.port == 8080

        # For ModularMCPServer, we can't easily test tool registration without mocking
        # The tools would be registered through the registry system
        assert server.tool_registry is not None

    def test_create_knowledge_base_server(self):
        """Test creating a knowledge base server"""
        kb_data = {
            "policy": "Company vacation policy",
            "faq": "Frequently asked questions"
        }

        server = MCPServerFactory.create_knowledge_base_server(
            name="KB Server",
            kb_data=kb_data
        )

        assert isinstance(server, ModularMCPServer)
        assert server.name == "KB Server"

        # For ModularMCPServer, we can't easily test tool registration without mocking
        # The tools would be registered through the registry system
        assert server.tool_registry is not None


class TestModularMCPServer:
    """Test ModularMCP server implementation"""

    @pytest.fixture
    def modular_server(self):
        """Create a modular server for testing"""
        return ModularMCPServer(
            name="Test Modular Server",
            host="localhost",
            port=8051
        )

    @pytest.mark.asyncio
    async def test_modular_server_initialization(self, modular_server):
        """Test modular server initialization"""
        assert modular_server.name == "Test Modular Server"
        assert modular_server.host == "localhost"
        assert modular_server.port == 8051
        assert not modular_server._initialized
        assert modular_server.tool_registry is not None
        assert modular_server.prompt_registry is not None
        assert modular_server.resource_registry is not None

    @pytest.mark.asyncio
    async def test_modular_server_initialize(self, modular_server):
        """Test modular server initialization process"""
        # Mock the registry methods to avoid actual file discovery
        with patch.object(modular_server.tool_registry, 'register_tools_with_server') as mock_tools, \
             patch.object(modular_server.tool_registry, 'get_all_tools', return_value=[]), \
             patch.object(modular_server.prompt_registry, 'register_prompts_with_server') as mock_prompts, \
             patch.object(modular_server.prompt_registry, 'get_all_prompts', return_value=[]), \
             patch.object(modular_server.resource_registry, 'register_resources_with_server') as mock_resources, \
             patch.object(modular_server.resource_registry, 'get_all_resources', return_value=[]):

            await modular_server.initialize()

            assert modular_server._initialized
            mock_tools.assert_called_once_with(modular_server.mcp)
            mock_prompts.assert_called_once_with(modular_server.mcp)
            mock_resources.assert_called_once_with(modular_server.mcp)

    @pytest.mark.asyncio
    async def test_modular_server_double_initialize(self, modular_server):
        """Test that double initialization doesn't cause issues"""
        with patch.object(modular_server.tool_registry, 'register_tools_with_server') as mock_tools, \
             patch.object(modular_server.tool_registry, 'get_all_tools', return_value=[]), \
             patch.object(modular_server.prompt_registry, 'register_prompts_with_server') as mock_prompts, \
             patch.object(modular_server.prompt_registry, 'get_all_prompts', return_value=[]), \
             patch.object(modular_server.resource_registry, 'register_resources_with_server') as mock_resources, \
             patch.object(modular_server.resource_registry, 'get_all_resources', return_value=[]):

            # Initialize twice
            await modular_server.initialize()
            await modular_server.initialize()

            # Should only be called once due to _initialized check
            mock_tools.assert_called_once()
            mock_prompts.assert_called_once()
            mock_resources.assert_called_once()

    def test_modular_server_get_server_info(self, modular_server):
        """Test getting server information"""
        # Mock the registry methods
        with patch.object(modular_server.tool_registry, 'get_all_tools', return_value=[Mock(name="tool1")]), \
             patch.object(modular_server.tool_registry, 'get_tool_names', return_value=["tool1"]), \
             patch.object(modular_server.prompt_registry, 'get_all_prompts', return_value=[Mock(name="prompt1")]), \
             patch.object(modular_server.prompt_registry, 'get_prompt_names', return_value=["prompt1"]), \
             patch.object(modular_server.resource_registry, 'get_all_resources', return_value=[Mock(uri="resource1")]), \
             patch.object(modular_server.resource_registry, 'get_resource_uris', return_value=["resource1"]):

            info = modular_server.get_server_info()

            assert info["name"] == "Test Modular Server"
            assert info["host"] == "localhost"
            assert info["port"] == 8051
            assert info["tools"]["count"] == 1
            assert info["tools"]["names"] == ["tool1"]
            assert info["prompts"]["count"] == 1
            assert info["prompts"]["names"] == ["prompt1"]
            assert info["resources"]["count"] == 1
            assert info["resources"]["uris"] == ["resource1"]

    def test_modular_server_run_without_initialize(self, modular_server):
        """Test that run method initializes server if not already initialized"""
        with patch.object(modular_server, 'initialize') as mock_init, \
             patch.object(modular_server.mcp, 'run') as mock_run:

            modular_server.run("stdio")

            mock_init.assert_called_once()
            mock_run.assert_called_once_with(transport="stdio")

    def test_modular_server_run_already_initialized(self, modular_server):
        """Test that run method doesn't reinitialize if already initialized"""
        modular_server._initialized = True

        with patch.object(modular_server, 'initialize') as mock_init, \
             patch.object(modular_server.mcp, 'run') as mock_run:

            modular_server.run("sse")

            mock_init.assert_not_called()
            mock_run.assert_called_once_with(transport="sse")

    def test_modular_server_custom_directories(self):
        """Test modular server with custom directories"""
        server = ModularMCPServer(
            name="Custom Server",
            tools_directory="/custom/tools",
            prompts_directory="/custom/prompts",
            resources_directory="/custom/resources"
        )

        assert server.tool_registry.directory == "/custom/tools"
        assert server.prompt_registry.directory == "/custom/prompts"
        assert server.resource_registry.directory == "/custom/resources"


class TestServerRegistries:
    """Test server registry classes"""

    @pytest.fixture
    def mock_tool(self):
        """Create a mock tool for testing"""
        tool = Mock()
        tool.name = "test_tool"
        tool.description = "A test tool"
        tool.input_schema = {"type": "object"}
        return tool

    @pytest.fixture
    def mock_prompt(self):
        """Create a mock prompt for testing"""
        prompt = Mock()
        prompt.name = "test_prompt"
        prompt.description = "A test prompt"
        prompt.arguments = {"name": {"type": "string"}}
        return prompt

    @pytest.fixture
    def mock_resource(self):
        """Create a mock resource for testing"""
        resource = Mock()
        resource.uri = "test://resource"
        resource.name = "Test Resource"
        resource.description = "A test resource"
        return resource

    def test_tool_registry_initialization(self):
        """Test tool registry initialization"""
        registry = ServerToolRegistry("/test/tools")
        assert registry.directory == "/test/tools"
        assert registry._tools == {}

    def test_tool_registry_register_tool(self, mock_tool):
        """Test tool registration"""
        registry = ServerToolRegistry()
        registry.register_tool(mock_tool)

        assert "test_tool" in registry._tools
        assert registry._tools["test_tool"] == mock_tool

    def test_tool_registry_get_all_tools(self, mock_tool):
        """Test getting all tools"""
        registry = ServerToolRegistry()
        registry.register_tool(mock_tool)

        tools = registry.get_all_tools()
        assert len(tools) == 1
        assert tools[0] == mock_tool

    def test_tool_registry_get_tool_names(self, mock_tool):
        """Test getting tool names"""
        registry = ServerToolRegistry()
        registry.register_tool(mock_tool)

        names = registry.get_tool_names()
        assert names == ["test_tool"]

    def test_prompt_registry_initialization(self):
        """Test prompt registry initialization"""
        registry = ServerPromptRegistry("/test/prompts")
        assert registry.directory == "/test/prompts"
        assert registry._prompts == {}

    def test_prompt_registry_register_prompt(self, mock_prompt):
        """Test prompt registration"""
        registry = ServerPromptRegistry()
        registry.register_prompt(mock_prompt)

        assert "test_prompt" in registry._prompts
        assert registry._prompts["test_prompt"] == mock_prompt

    def test_prompt_registry_get_all_prompts(self, mock_prompt):
        """Test getting all prompts"""
        registry = ServerPromptRegistry()
        registry.register_prompt(mock_prompt)

        prompts = registry.get_all_prompts()
        assert len(prompts) == 1
        assert prompts[0] == mock_prompt

    def test_prompt_registry_get_prompt_names(self, mock_prompt):
        """Test getting prompt names"""
        registry = ServerPromptRegistry()
        registry.register_prompt(mock_prompt)

        names = registry.get_prompt_names()
        assert names == ["test_prompt"]

    def test_resource_registry_initialization(self):
        """Test resource registry initialization"""
        registry = ServerResourceRegistry("/test/resources")
        assert registry.directory == "/test/resources"
        assert registry._resources == {}

    def test_resource_registry_register_resource(self, mock_resource):
        """Test resource registration"""
        registry = ServerResourceRegistry()
        registry.register_resource(mock_resource)

        assert "test://resource" in registry._resources
        assert registry._resources["test://resource"] == mock_resource

    def test_resource_registry_get_all_resources(self, mock_resource):
        """Test getting all resources"""
        registry = ServerResourceRegistry()
        registry.register_resource(mock_resource)

        resources = registry.get_all_resources()
        assert len(resources) == 1
        assert resources[0] == mock_resource

    def test_resource_registry_get_resource_uris(self, mock_resource):
        """Test getting resource URIs"""
        registry = ServerResourceRegistry()
        registry.register_resource(mock_resource)

        uris = registry.get_resource_uris()
        assert uris == ["test://resource"]


class TestBaseServerTools:
    """Test base server tool classes"""

    def test_base_server_tool_initialization(self):
        """Test BaseServerTool initialization"""
        class TestTool(BaseServerTool):
            async def execute(self, **kwargs):
                return "test_result"

        tool = TestTool(
            name="test_tool",
            description="A test tool",
            input_schema={"type": "object"}
        )

        assert tool.name == "test_tool"
        assert tool.description == "A test tool"
        assert tool.input_schema == {"type": "object"}

    def test_base_server_tool_get_tool_definition(self):
        """Test getting tool definition"""
        class TestTool(BaseServerTool):
            async def execute(self, **kwargs):
                return "test_result"

        tool = TestTool(
            name="test_tool",
            description="A test tool",
            input_schema={"type": "object"}
        )

        definition = tool.get_tool_definition()
        expected = {
            "name": "test_tool",
            "description": "A test tool",
            "input_schema": {"type": "object"}
        }

        assert definition == expected

    def test_base_server_tool_create_fastmcp_tool(self):
        """Test creating FastMCP tool wrapper"""
        class TestTool(BaseServerTool):
            async def execute(self, **kwargs):
                return "test_result"

        tool = TestTool(
            name="test_tool",
            description="A test tool",
            input_schema={"type": "object"}
        )

        # Mock MCP server
        mock_mcp = Mock()
        mock_mcp.tool.return_value = lambda func: func

        wrapper = tool.create_fastmcp_tool(mock_mcp)

        assert wrapper.__name__ == "test_tool"
        assert wrapper.__doc__ == "A test tool"

    def test_context_aware_tool_initialization(self):
        """Test ContextAwareTool initialization"""
        class TestContextTool(ContextAwareTool):
            async def execute_with_context(self, ctx, **kwargs):
                return "context_result"

        tool = TestContextTool(
            name="context_tool",
            description="A context-aware tool",
            input_schema={"type": "object"}
        )

        assert tool.name == "context_tool"
        assert tool.description == "A context-aware tool"
        assert tool.input_schema == {"type": "object"}

    @pytest.mark.asyncio
    async def test_context_aware_tool_execute_with_context(self):
        """Test ContextAwareTool execution with context"""
        from unittest.mock import Mock

        class TestContextTool(ContextAwareTool):
            async def execute_with_context(self, ctx, **kwargs):
                return f"Context: {self.context}, Args: {kwargs}"

        tool = TestContextTool(
            name="context_tool",
            description="A context-aware tool",
            input_schema={"type": "object"}
        )

        # Set context
        tool.context = {"user_id": "123"}

        # Mock the context
        mock_ctx = Mock()

        result = await tool.execute_with_context(mock_ctx, test_arg="value")
        assert "Context: {'user_id': '123'}" in result
        assert "Args: {'test_arg': 'value'}" in result