python-sdk/pkg/hanzo-memory/tests/test_mcp_server.py

"""Tests for MCP server."""

from unittest.mock import MagicMock, patch

import pytest

from hanzo_memory.mcp.server import MCPMemoryServer


@pytest.fixture
def mcp_server():
    """Create MCP server instance."""
    with (
        patch("hanzo_memory.mcp.server.get_db_client"),
        patch("hanzo_memory.mcp.server.EmbeddingService"),
        patch("hanzo_memory.mcp.server.LLMService"),
    ):
        server = MCPMemoryServer()
        # Mock the embedding service
        server.embedding_service.embed_text = MagicMock(return_value=[[0.1] * 384])
        return server


@pytest.mark.asyncio
async def test_tools_are_registered(mcp_server):
    """Test that tools are properly registered."""
    # Just verify that _setup_handlers was called and tools exist
    # The actual tool listing would be tested during integration
    assert hasattr(mcp_server, "server")
    assert hasattr(mcp_server, "db_client")
    assert hasattr(mcp_server, "embedding_service")
    assert hasattr(mcp_server, "llm_service")


@pytest.mark.asyncio
async def test_handle_remember(mcp_server):
    """Test remember tool."""
    # Mock DB operations
    mcp_server.db_client.create_memories_table = MagicMock()
    mcp_server.db_client.add_memory = MagicMock(return_value={"memory_id": "test-id"})

    result = await mcp_server._handle_remember(
        {
            "user_id": "user123",
            "project_id": "proj123",
            "content": "Test memory content",
            "metadata": {"tag": "test"},
            "importance": 5.0,
        }
    )

    assert result["success"] is True
    assert "memory_id" in result
    assert result["message"] == "Memory stored successfully"

    # Verify DB calls
    mcp_server.db_client.create_memories_table.assert_called_once_with("user123")
    mcp_server.db_client.add_memory.assert_called_once()


@pytest.mark.asyncio
async def test_handle_recall(mcp_server):
    """Test recall tool."""
    # Mock search results
    import polars as pl

    mock_df = pl.DataFrame(
        {
            "memory_id": ["mem1", "mem2"],
            "content": ["Memory 1", "Memory 2"],
            "metadata": ['{"tag": "test1"}', '{"tag": "test2"}'],
            "importance": [5.0, 3.0],
            "_similarity": [0.9, 0.7],
        }
    )

    mcp_server.db_client.search_memories = MagicMock(return_value=mock_df)

    result = await mcp_server._handle_recall(
        {
            "user_id": "user123",
            "query": "test query",
            "limit": 5,
        }
    )

    assert result["success"] is True
    assert len(result["memories"]) == 2
    assert result["memories"][0]["content"] == "Memory 1"
    assert result["memories"][0]["similarity_score"] == 0.9


@pytest.mark.asyncio
async def test_handle_create_project(mcp_server):
    """Test create project tool."""
    mcp_server.db_client.create_project = MagicMock(
        return_value={"project_id": "proj123"}
    )

    result = await mcp_server._handle_create_project(
        {
            "user_id": "user123",
            "name": "Test Project",
            "description": "A test project",
            "metadata": {"category": "test"},
        }
    )

    assert result["success"] is True
    assert "project_id" in result
    assert result["message"] == "Project created successfully"


@pytest.mark.asyncio
async def test_handle_create_knowledge_base(mcp_server):
    """Test create knowledge base tool."""
    mcp_server.db_client.create_knowledge_base = MagicMock(
        return_value={"kb_id": "kb123"}
    )

    result = await mcp_server._handle_create_knowledge_base(
        {
            "user_id": "user123",
            "project_id": "proj123",
            "name": "Test KB",
            "description": "A test knowledge base",
        }
    )

    assert result["success"] is True
    assert "kb_id" in result
    assert result["message"] == "Knowledge base created successfully"


@pytest.mark.asyncio
async def test_handle_add_fact(mcp_server):
    """Test add fact tool."""
    mcp_server.db_client.add_fact = MagicMock(return_value={"fact_id": "fact123"})

    result = await mcp_server._handle_add_fact(
        {
            "kb_id": "kb123",
            "content": "Test fact content",
            "parent_id": "parent123",
            "metadata": {"type": "definition"},
        }
    )

    assert result["success"] is True
    assert "fact_id" in result
    assert result["message"] == "Fact added successfully"


@pytest.mark.asyncio
async def test_handle_search_facts(mcp_server):
    """Test search facts tool."""
    import polars as pl

    mock_df = pl.DataFrame(
        {
            "fact_id": ["fact1", "fact2"],
            "content": ["Fact 1", "Fact 2"],
            "parent_id": ["", "parent1"],
            "metadata": ['{"type": "def"}', '{"type": "example"}'],
            "_similarity": [0.95, 0.85],
        }
    )

    mcp_server.db_client.search_facts = MagicMock(return_value=mock_df)

    result = await mcp_server._handle_search_facts(
        {
            "kb_id": "kb123",
            "query": "test query",
            "limit": 10,
        }
    )

    assert result["success"] is True
    assert len(result["facts"]) == 2
    assert result["facts"][0]["content"] == "Fact 1"
    assert result["facts"][0]["similarity_score"] == 0.95


@pytest.mark.asyncio
async def test_handle_summarize_for_knowledge(mcp_server):
    """Test summarize for knowledge tool."""
    mcp_server.llm_service.summarize_for_knowledge = MagicMock(
        return_value={
            "summary": "Test summary",
            "knowledge_instructions": {
                "action": "add_fact",
                "facts": [{"content": "Extracted fact"}],
                "reasoning": "Test reasoning",
            },
        }
    )

    result = await mcp_server._handle_summarize_for_knowledge(
        {
            "content": "Long content to summarize",
            "context": "Additional context",
            "skip_summarization": False,
        }
    )

    assert "summary" in result
    assert "knowledge_instructions" in result
    assert result["knowledge_instructions"]["action"] == "add_fact"


@pytest.mark.asyncio
async def test_handle_tool_error_handling(mcp_server):
    """Test error handling in tool calls."""
    # Mock error
    mcp_server.db_client.create_memories_table = MagicMock(
        side_effect=Exception("DB Error")
    )

    # Call the handler directly instead of through decorated function
    try:
        await mcp_server._handle_remember(
            {
                "user_id": "user123",
                "project_id": "proj123",
                "content": "Test memory",
            }
        )
        raise AssertionError("Should have raised an exception")
    except Exception as e:
        assert "DB Error" in str(e)


@pytest.mark.asyncio
async def test_mcp_server_initialization(mcp_server):
    """Test MCP server is properly initialized."""
    # Verify the server is set up correctly
    assert mcp_server.server is not None
    assert mcp_server.server.name == "hanzo-memory"

    # Verify services are initialized
    assert mcp_server.db_client is not None
    assert mcp_server.embedding_service is not None
    assert mcp_server.llm_service is not None

    # The actual MCP protocol testing would be done via integration tests
    # with the full MCP framework running