nocobase-mcp/tests/test_server.py

import asyncio
import unittest
from unittest.mock import patch

import httpx
from starlette.responses import JSONResponse

from nocobase_crm_mcp.server import (
    DEFAULT_HTTP_HOST,
    DEFAULT_HTTP_PATH,
    DEFAULT_HTTP_PORT,
    BearerTokenAuthMiddleware,
    FastMCP,
    create_http_app,
    create_server,
    run_http_server,
)


async def _ok_app(scope, receive, send):
    await send({"type": "http.response.start", "status": 200, "headers": []})
    await send({"type": "http.response.body", "body": b"ok"})


def run_asgi_request(app, *, authorization=None, method="POST", path="/mcp"):
    events = []
    headers = []
    if authorization is not None:
        headers.append((b"authorization", authorization.encode("latin1")))

    async def receive():
        return {"type": "http.request", "body": b"", "more_body": False}

    async def send(message):
        events.append(message)

    scope = {"type": "http", "method": method, "path": path, "headers": headers}
    asyncio.run(app(scope, receive, send))
    return events


class FakeCrmService:
    def crm_health(self):
        return {"ok": True}

    def crm_schema(self, include_fields=False):
        return {"include_fields": include_fields}

    def crm_overview(self, **kwargs):
        return {"totals": {"customers": 1}}

    def crm_config_template(self):
        return {"collections": {}}

    def crm_filter_options(self, **kwargs):
        return {"customer_status": []}

    def resolve_owner(self, **kwargs):
        return {"results": []}

    def list_customers(self, **kwargs):
        return {"results": []}

    def get_customer(self, **kwargs):
        return {"customer": None}

    def list_contacts(self, **kwargs):
        return {"results": []}

    def list_followups(self, **kwargs):
        return {"results": []}


class FakeProjectService:
    def create_project(self, **kwargs):
        return {"project_id": 1}

    def update_project(self, project_id, **kwargs):
        return {"project_id": project_id}


class FakeUploadEndpoint:
    async def __call__(self, request):
        return JSONResponse({"uploaded": True})


class ServerTests(unittest.TestCase):
    def test_bearer_token_middleware_rejects_missing_or_wrong_token(self) -> None:
        app = BearerTokenAuthMiddleware(_ok_app, "secret")

        self.assertEqual(run_asgi_request(app)[0]["status"], 401)
        self.assertEqual(run_asgi_request(app, authorization="Bearer wrong")[0]["status"], 401)

    def test_bearer_token_middleware_allows_matching_token(self) -> None:
        app = BearerTokenAuthMiddleware(_ok_app, "secret")

        events = run_asgi_request(app, authorization="Bearer secret")

        self.assertEqual(events[0]["status"], 200)
        self.assertEqual(events[1]["body"], b"ok")

    def test_bearer_token_middleware_allows_options_preflight_without_token(self) -> None:
        app = BearerTokenAuthMiddleware(_ok_app, "secret")

        events = run_asgi_request(app, method="OPTIONS")

        self.assertEqual(events[0]["status"], 204)

    def test_readonly_server_exposes_crm_tools_only(self) -> None:
        if FastMCP is None:
            self.skipTest("MCP SDK is not installed")

        server = create_server(crm_service=FakeCrmService(), readonly=True)
        tools = asyncio.run(server.list_tools())

        self.assertEqual(
            {tool.name for tool in tools},
            {
                "crm_health",
                "crm_schema",
                "crm_overview",
                "crm_filter_options",
                "resolve_owner",
                "list_customers",
                "get_customer",
                "list_contacts",
                "list_followups",
            },
        )

    def test_write_server_exposes_fixed_project_tools(self) -> None:
        if FastMCP is None:
            self.skipTest("MCP SDK is not installed")

        server = create_server(
            crm_service=FakeCrmService(),
            project_service=FakeProjectService(),
            readonly=False,
        )

        names = {tool.name for tool in asyncio.run(server.list_tools())}

        self.assertTrue({"create_project", "update_project"}.issubset(names))
        self.assertNotIn("destroy_project", names)

    def test_upload_route_requires_the_same_bearer_token(self) -> None:
        if FastMCP is None:
            self.skipTest("MCP SDK is not installed")

        server = create_server(crm_service=FakeCrmService(), readonly=True)
        app = create_http_app(server, "secret", upload_endpoint=FakeUploadEndpoint())

        self.assertEqual(run_asgi_request(app, path="/uploads")[0]["status"], 401)

    def test_upload_route_accepts_the_matching_bearer_token(self) -> None:
        if FastMCP is None:
            self.skipTest("MCP SDK is not installed")

        server = create_server(crm_service=FakeCrmService(), readonly=True)

        async def upload_endpoint(request):
            return JSONResponse({"uploaded": True})

        app = create_http_app(server, "secret", upload_endpoint=upload_endpoint)

        async def request_upload() -> httpx.Response:
            async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://testserver") as client:
                return await client.post("/uploads", headers={"Authorization": "Bearer secret"})

        response = asyncio.run(request_upload())
        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json(), {"uploaded": True})

    def test_http_write_mode_requires_upload_api_base_url(self) -> None:
        with patch.dict("os.environ", {"NOCOBASE_MCP_TOKEN": "secret"}, clear=True):
            with patch("uvicorn.run") as uvicorn_run:
                with self.assertRaisesRegex(RuntimeError, "NOCOBASE_MCP_UPLOAD_API_BASE_URL"):
                    run_http_server(
                        config_path=None,
                        host=DEFAULT_HTTP_HOST,
                        port=DEFAULT_HTTP_PORT,
                        token_env="NOCOBASE_MCP_TOKEN",
                        streamable_http_path=DEFAULT_HTTP_PATH,
                        allowed_hosts=None,
                        allowed_origins=None,
                    )
                uvicorn_run.assert_not_called()

    def test_config_template_tool_is_opt_in(self) -> None:
        if FastMCP is None:
            self.skipTest("MCP SDK is not installed")

        with patch.dict("os.environ", {"NOCOBASE_MCP_ENABLE_CONFIG_TOOLS": "1"}):
            server = create_server(crm_service=FakeCrmService(), readonly=True)
        tools = asyncio.run(server.list_tools())

        self.assertIn("crm_config_template", {tool.name for tool in tools})

    def test_create_http_app_requires_bearer_token(self) -> None:
        if FastMCP is None:
            self.skipTest("MCP SDK is not installed")

        server = create_server(crm_service=FakeCrmService(), readonly=True)
        with self.assertRaises(ValueError):
            create_http_app(server, "")


if __name__ == "__main__":
    unittest.main()