Memory MCP Server

MIT License
OverviewInspectSchema Related Servers Reviews Score
tests
"""Tests for the MCP server implementation."""

import json
from typing import Any, Dict, List, Protocol, cast

import pytest
from mcp.types import TextContent

from memory_mcp_server.exceptions import EntityNotFoundError
from memory_mcp_server.interfaces import Entity, KnowledgeGraph, Relation


# Mock tools and handlers
def handle_error(error: Exception) -> str:
    """Mock error handler."""
    if isinstance(error, EntityNotFoundError):
        return str(error)
    return f"Error: {str(error)}"


async def create_entities_handler(
    manager: Any, arguments: Dict[str, Any]
) -> List[TextContent]:
    """Mock create entities handler."""
    entities = [
        Entity(
            name=e["name"],
            entityType=e["entityType"],
            observations=e.get("observations", []),
        )
        for e in arguments["entities"]
    ]
    result = await manager.create_entities(entities)
    return [TextContent(type="text", text=json.dumps([e.to_dict() for e in result]))]


async def create_relations_handler(
    manager: Any, arguments: Dict[str, Any]
) -> List[TextContent]:
    """Mock create relations handler."""
    relations = [
        Relation(from_=r["from"], to=r["to"], relationType=r["relationType"])
        for r in arguments["relations"]
    ]
    result = await manager.create_relations(relations)
    return [TextContent(type="text", text=json.dumps([r.to_dict() for r in result]))]


async def add_observations_handler(
    manager: Any, arguments: Dict[str, Any]
) -> List[TextContent]:
    """Mock add observations handler."""
    await manager.add_observations(arguments["entity"], arguments["observations"])
    return [TextContent(type="text", text=json.dumps({"success": True}))]


async def delete_entities_handler(
    manager: Any, arguments: Dict[str, Any]
) -> List[TextContent]:
    """Mock delete entities handler."""
    await manager.delete_entities(arguments["names"])
    return [TextContent(type="text", text=json.dumps({"success": True}))]


async def delete_observations_handler(
    manager: Any, arguments: Dict[str, Any]
) -> List[TextContent]:
    """Mock delete observations handler."""
    await manager.delete_observations(arguments["entity"], arguments["observations"])
    return [TextContent(type="text", text=json.dumps({"success": True}))]


async def delete_relations_handler(
    manager: Any, arguments: Dict[str, Any]
) -> List[TextContent]:
    """Mock delete relations handler."""
    await manager.delete_relations(arguments["from"], arguments["to"])
    return [TextContent(type="text", text=json.dumps({"success": True}))]


async def read_graph_handler(
    manager: Any, arguments: Dict[str, Any]
) -> List[TextContent]:
    """Mock read graph handler."""
    graph = await manager.read_graph()
    return [TextContent(type="text", text=json.dumps(graph.to_dict()))]


async def search_nodes_handler(
    manager: Any, arguments: Dict[str, Any]
) -> List[TextContent]:
    """Mock search nodes handler."""
    result = await manager.search_nodes(arguments["query"])
    return [TextContent(type="text", text=json.dumps(result.to_dict()))]


TOOLS: Dict[str, Any] = {
    "create_entities": create_entities_handler,
    "create_relations": create_relations_handler,
    "add_observations": add_observations_handler,
    "delete_entities": delete_entities_handler,
    "delete_relations": delete_relations_handler,
    "read_graph": read_graph_handler,
    "search_nodes": search_nodes_handler,
}


class MockManagerProtocol(Protocol):
    """Protocol defining the interface for MockManager."""

    async def create_entities(self, entities: List[Entity]) -> List[Entity]:
        ...

    async def create_relations(self, relations: List[Relation]) -> List[Relation]:
        ...

    async def add_observations(self, entity: str, observations: List[str]) -> None:
        ...

    async def delete_entities(self, names: List[str]) -> None:
        ...

    async def delete_relations(self, from_: str, to: str) -> None:
        ...

    async def read_graph(self) -> KnowledgeGraph:
        ...

    async def search_nodes(self, query: str) -> KnowledgeGraph:
        ...

    async def flush(self) -> None:
        ...


@pytest.fixture(scope="function")
def mock_manager() -> MockManagerProtocol:
    """Create a mock manager for testing."""

    class MockManager:
        def __init__(self) -> None:
            self.entities: List[Entity] = []
            self.relations: List[Relation] = []

        async def create_entities(self, entities: List[Entity]) -> List[Entity]:
            self.entities.extend(entities)
            return entities

        async def create_relations(self, relations: List[Relation]) -> List[Relation]:
            return relations

        async def add_observations(self, entity: str, observations: List[str]) -> None:
            if entity == "MissingEntity":
                raise EntityNotFoundError(entity)

        async def delete_entities(self, names: List[str]) -> None:
            for name in names:
                if name == "MissingEntity":
                    raise EntityNotFoundError(name)

        async def delete_relations(self, from_: str, to: str) -> None:
            if from_ == "MissingEntity" or to == "MissingEntity":
                raise EntityNotFoundError("MissingEntity")

        async def read_graph(self) -> KnowledgeGraph:
            # Return current state including any created entities
            return KnowledgeGraph(
                entities=self.entities,
                relations=self.relations,
            )

        async def search_nodes(self, query: str) -> KnowledgeGraph:
            # If query matches "TestEntity", return graph; otherwise empty
            if "TestEntity".lower() in query.lower():
                return await self.read_graph()
            return KnowledgeGraph(entities=[], relations=[])

        async def open_nodes(self, names: List[str]) -> KnowledgeGraph:
            # If "TestEntity" is requested, return it
            if "TestEntity" in names:
                return await self.read_graph()
            return KnowledgeGraph(entities=[], relations=[])

        async def flush(self) -> None:
            """Mock flush method to comply with interface."""
            pass

    return MockManager()


@pytest.mark.asyncio
async def test_create_entities(mock_manager: MockManagerProtocol) -> None:
    """Test creating entities through the MCP server."""
    handler = cast(Any, TOOLS["create_entities"])
    arguments = {
        "entities": [
            {"name": "E1", "entityType": "TypeX", "observations": ["obsA"]},
            {"name": "E2", "entityType": "TypeY", "observations": ["obsB"]},
        ]
    }
    result = await handler(mock_manager, arguments)
    data = json.loads(result[0].text)
    assert len(data) == 2
    assert data[0]["name"] == "E1"
    assert data[1]["observations"] == ["obsB"]


@pytest.mark.asyncio
async def test_create_relations(mock_manager: MockManagerProtocol) -> None:
    """Test creating relations through the MCP server."""
    handler = cast(Any, TOOLS["create_relations"])
    arguments = {"relations": [{"from": "E1", "to": "E2", "relationType": "likes"}]}
    result = await handler(mock_manager, arguments)
    data = json.loads(result[0].text)
    assert len(data) == 1
    assert data[0]["from"] == "E1"
    assert data[0]["to"] == "E2"


@pytest.mark.asyncio
async def test_add_observations(mock_manager: MockManagerProtocol) -> None:
    """Test adding observations through the MCP server."""
    handler = cast(Any, TOOLS["add_observations"])
    arguments = {"entity": "E1", "observations": ["newObs"]}
    result = await handler(mock_manager, arguments)
    data = json.loads(result[0].text)
    assert data["success"] is True


@pytest.mark.asyncio
async def test_delete_entities(mock_manager: MockManagerProtocol) -> None:
    """Test deleting entities through the MCP server."""
    handler = cast(Any, TOOLS["delete_entities"])
    arguments = {"names": ["E1"]}
    result = await handler(mock_manager, arguments)
    data = json.loads(result[0].text)
    assert data["success"] is True


@pytest.mark.asyncio
async def test_delete_relations(mock_manager: MockManagerProtocol) -> None:
    """Test deleting relations through the MCP server."""
    handler = cast(Any, TOOLS["delete_relations"])
    arguments = {"from": "E1", "to": "E2"}
    result = await handler(mock_manager, arguments)
    data = json.loads(result[0].text)
    assert data["success"] is True


@pytest.mark.asyncio
async def test_read_graph(mock_manager: MockManagerProtocol) -> None:
    """Test reading the graph through the MCP server."""
    # Create test entity first
    await mock_manager.create_entities(
        [
            Entity(
                name="TestEntity",
                entityType="TestType",
                observations=["test observation"],
            )
        ]
    )

    handler = cast(Any, TOOLS["read_graph"])
    arguments: Dict[str, Any] = {}
    result = await handler(mock_manager, arguments)
    data = json.loads(result[0].text)

    assert len(data["entities"]) == 1
    assert data["entities"][0]["name"] == "TestEntity"
    assert isinstance(
        data["entities"][0]["observations"], (list, tuple)
    )  # Allow both list and tuple


@pytest.mark.asyncio
async def test_save_graph(mock_manager: MockManagerProtocol) -> None:
    """Test saving the graph through the MCP server."""
    # First create some test data
    await mock_manager.create_entities(
        [Entity(name="TestSave", entityType="TestType", observations=[])]
    )

    # Explicitly save the graph
    await mock_manager.flush()

    # Read back the graph
    graph = await mock_manager.read_graph()

    # Verify our test entity exists
    assert any(e.name == "TestSave" for e in graph.entities)
    # Verify the save preserved the structure
    assert isinstance(
        graph.entities[0].observations, (list, tuple)
    )  # Allow both list and tuple for immutability


@pytest.mark.asyncio
async def test_search_nodes(mock_manager: MockManagerProtocol) -> None:
    """Test searching nodes through the MCP server."""
    # Create test entity first
    await mock_manager.create_entities(
        [
            Entity(
                name="TestEntity",
                entityType="TestType",
                observations=["test observation"],
            )
        ]
    )

    handler = cast(Any, TOOLS["search_nodes"])
    arguments = {"query": "TestEntity"}
    result = await handler(mock_manager, arguments)
    data = json.loads(result[0].text)

    assert len(data["entities"]) == 1
    assert data["entities"][0]["name"] == "TestEntity"
    assert isinstance(
        data["entities"][0]["observations"], (list, tuple)
    )  # Allow both list and tuple


def test_error_handling() -> None:
    """Test error handling functionality."""
    msg = handle_error(EntityNotFoundError("MissingEntity"))
    assert "Entity 'MissingEntity' not found in the graph" in msg