tests/test_mcp_server.py

import sys
import types

import pytest

from mnemocore.core.config import HAIMConfig, MCPConfig, SecurityConfig
from mnemocore.mcp.adapters.api_adapter import MnemoCoreAPIError
from mnemocore.mcp import server as mcp_server


class FakeFastMCP:
    def __init__(self, name: str):
        self.name = name
        self.tools = {}
        self.run_calls = []

    def tool(self):
        def decorator(fn):
            self.tools[fn.__name__] = fn
            return fn

        return decorator

    def run(self, **kwargs):
        self.run_calls.append(kwargs)


class FakeAdapter:
    def store(self, payload):
        return {"memory_id": "mem_1", "payload": payload}

    def query(self, payload):
        return {"results": [{"id": "mem_1", "score": 1.0}], "payload": payload}

    def get_memory(self, memory_id: str):
        return {"id": memory_id, "content": "hello"}

    def delete_memory(self, memory_id: str):
        return {"deleted": memory_id}

    def stats(self):
        return {"engine_version": "3.5.1"}

    def health(self):
        return {"status": "healthy"}


def _install_fake_mcp_modules(monkeypatch):
    mcp_mod = types.ModuleType("mcp")
    server_mod = types.ModuleType("mcp.server")
    fastmcp_mod = types.ModuleType("mcp.server.fastmcp")
    fastmcp_mod.FastMCP = FakeFastMCP

    monkeypatch.setitem(sys.modules, "mcp", mcp_mod)
    monkeypatch.setitem(sys.modules, "mcp.server", server_mod)
    monkeypatch.setitem(sys.modules, "mcp.server.fastmcp", fastmcp_mod)


def test_build_server_registers_only_allowlisted_tools(monkeypatch):
    _install_fake_mcp_modules(monkeypatch)
    monkeypatch.setattr(mcp_server, "MnemoCoreAPIAdapter", lambda *args, **kwargs: FakeAdapter())

    config = HAIMConfig(
        security=SecurityConfig(api_key="test-key"),
        mcp=MCPConfig(
            enabled=True,
            allow_tools=["memory_health", "memory_stats"],
            api_key="test-key",
        ),
    )

    server = mcp_server.build_server(config)
    assert sorted(server.tools.keys()) == ["memory_health", "memory_stats"]

    health_result = server.tools["memory_health"]()
    assert health_result["ok"] is True
    assert health_result["data"]["status"] == "healthy"


def test_tool_error_handling(monkeypatch):
    class ErrorAdapter(FakeAdapter):
        def health(self):
            raise MnemoCoreAPIError("boom", status_code=503)

    _install_fake_mcp_modules(monkeypatch)
    monkeypatch.setattr(mcp_server, "MnemoCoreAPIAdapter", lambda *args, **kwargs: ErrorAdapter())

    config = HAIMConfig(
        security=SecurityConfig(api_key="test-key"),
        mcp=MCPConfig(enabled=True, allow_tools=["memory_health"], api_key="test-key"),
    )

    server = mcp_server.build_server(config)
    result = server.tools["memory_health"]()

    assert result["ok"] is False
    assert "boom" in result["error"]


def test_main_runs_with_stdio_transport(monkeypatch):
    fake_server = FakeFastMCP("x")

    monkeypatch.setattr(
        mcp_server,
        "get_config",
        lambda: HAIMConfig(
            security=SecurityConfig(api_key="k"),
            mcp=MCPConfig(enabled=True, transport="stdio", api_key="k"),
        ),
    )
    monkeypatch.setattr(mcp_server, "build_server", lambda cfg: fake_server)

    mcp_server.main()
    assert fake_server.run_calls == [{"transport": "stdio"}]


def test_main_runs_with_sse_transport(monkeypatch):
    fake_server = FakeFastMCP("x")

    monkeypatch.setattr(
        mcp_server,
        "get_config",
        lambda: HAIMConfig(
            security=SecurityConfig(api_key="k"),
            mcp=MCPConfig(
                enabled=True,
                transport="sse",
                host="127.0.0.1",
                port=8222,
                api_key="k",
            ),
        ),
    )
    monkeypatch.setattr(mcp_server, "build_server", lambda cfg: fake_server)

    mcp_server.main()
    assert fake_server.run_calls == [
        {"transport": "sse", "host": "127.0.0.1", "port": 8222}
    ]


def test_main_rejects_unknown_transport(monkeypatch):
    monkeypatch.setattr(
        mcp_server,
        "get_config",
        lambda: HAIMConfig(
            security=SecurityConfig(api_key="k"),
            mcp=MCPConfig(enabled=True, transport="unknown", api_key="k"),
        ),
    )
    monkeypatch.setattr(mcp_server, "build_server", lambda cfg: FakeFastMCP("x"))

    with pytest.raises((ValueError, Exception), match="Unsupported transport"):
        mcp_server.main()