agentic-soc-platform/backend/apps/mcp/tests.py

from asgiref.sync import async_to_sync
from django.test import SimpleTestCase, TestCase
from django.utils import timezone
from starlette.routing import Mount

from apps.accounts.models import User, UserApiKey
from apps.mcp.asgi import mcp_asgi_app
from apps.mcp.auth import ApiKeyMCPMiddleware
from apps.mcp.server import create_mcp_server
from apps.mcp.tools import REGISTERED_MCP_TOOLS


class FakeMCPServer:
    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs
        self.tools = []

    def add_tool(self, tool):
        self.tools.append(tool)


class MCPServerTests(SimpleTestCase):
    def test_create_mcp_server_registers_all_tools(self):
        server = create_mcp_server(server_factory=FakeMCPServer)

        self.assertEqual(server.args[0], "ASP-MCP")
        self.assertTrue(server.kwargs["stateless_http"])
        self.assertTrue(server.kwargs["json_response"])
        self.assertEqual(server.kwargs["streamable_http_path"], "/")
        self.assertEqual(server.tools, REGISTERED_MCP_TOOLS)

    def test_mcp_asgi_app_uses_api_key_middleware(self):
        self.assertIsInstance(mcp_asgi_app, ApiKeyMCPMiddleware)


class MCPASGIMountTests(SimpleTestCase):
    def test_asp_asgi_mounts_new_mcp_path_only(self):
        from asp.asgi import application

        http_application = application.application_mapping["http"]
        paths = [route.path for route in http_application.routes if isinstance(route, Mount)]

        self.assertIn("/api/mcp", paths)
        self.assertNotIn("/api/agentic/mcp", paths)


class MCPAuthMiddlewareTests(TestCase):
    def _call_middleware(self, middleware, *, headers=None):
        messages = []
        scope = {"type": "http", "headers": headers or []}

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

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

        async_to_sync(middleware)(scope, receive, send)
        return scope, messages

    def test_missing_api_key_returns_unauthorized(self):
        async def app(scope, receive, send):
            raise AssertionError("inner app should not be called")

        _scope, messages = self._call_middleware(ApiKeyMCPMiddleware(app))

        self.assertEqual(messages[0]["status"], 401)
        self.assertIn(b"API key required", messages[1]["body"])

    def test_invalid_api_key_returns_unauthorized(self):
        async def app(scope, receive, send):
            raise AssertionError("inner app should not be called")

        _scope, messages = self._call_middleware(
            ApiKeyMCPMiddleware(app),
            headers=[(b"authorization", b"Api-Key invalid")],
        )

        self.assertEqual(messages[0]["status"], 401)
        self.assertIn(b"Invalid API key", messages[1]["body"])

    def test_valid_api_key_sets_scope_and_updates_last_used_at(self):
        user = User.objects.create_user(username="mcp-user", password="unused")
        api_key = UserApiKey.objects.create(user=user, name="mcp", key="asp_test")
        captured_scope = {}

        async def app(scope, receive, send):
            captured_scope.update(scope)
            await send({"type": "http.response.start", "status": 204, "headers": []})
            await send({"type": "http.response.body", "body": b""})

        scope, messages = self._call_middleware(
            ApiKeyMCPMiddleware(app),
            headers=[(b"authorization", b"Api-Key asp_test")],
        )

        api_key.refresh_from_db()
        self.assertEqual(messages[0]["status"], 204)
        self.assertEqual(scope["user"], user)
        self.assertEqual(scope["api_key"], api_key)
        self.assertEqual(captured_scope["user"], user)
        self.assertIsNotNone(api_key.last_used_at)
        self.assertLessEqual(api_key.last_used_at, timezone.now())