ai-video-fullstack/backend/tests/test_mcp_client.py

import json
import unittest
from contextlib import asynccontextmanager
from unittest.mock import patch

import httpx2

from models import RuntimeMcpServer
from schemas import McpServerUpsert
from services.tools.mcp_client import McpClientError, McpToolClient


class McpTransportTests(unittest.IsolatedAsyncioTestCase):
    async def test_streamable_http_discovery_and_call_with_real_sdk(self):
        requests = []

        def respond(request):
            self.assertEqual(request.headers["Authorization"], "Bearer test")
            message = json.loads(request.content)
            requests.append(message)
            if "id" not in message:
                return httpx2.Response(202)
            if message["method"] == "initialize":
                result = {
                    "protocolVersion": "2025-06-18",
                    "capabilities": {"tools": {}},
                    "serverInfo": {"name": "test", "version": "1"},
                }
            elif message["method"] == "tools/list":
                if (message.get("params") or {}).get("cursor") == "page2":
                    result = {"tools": []}
                else:
                    result = {
                        "tools": [{
                            "name": "get_order",
                            "inputSchema": {"type": "object", "properties": {}},
                        }],
                        "nextCursor": "page2",
                    }
            elif message["method"] == "tools/call":
                result = {
                    "content": [{"type": "text", "text": "已发货"}],
                    "structuredContent": {"status": "shipped"},
                }
            else:
                self.fail(f"Unexpected MCP method: {message['method']}")
            return httpx2.Response(200, json={
                "jsonrpc": "2.0", "id": message["id"], "result": result,
            })

        client_class = httpx2.AsyncClient

        def client_factory(**kwargs):
            return client_class(transport=httpx2.MockTransport(respond), **kwargs)

        server = RuntimeMcpServer(
            id="mcp_http",
            url="https://mcp.example.com/mcp",
            secret_headers={"Authorization": "Bearer test"},
        )
        client = McpToolClient()
        with patch("services.tools.mcp_client.httpx.AsyncClient", client_factory):
            tools = await client.list_tools(server)
            listings = [item for item in requests if item["method"] == "tools/list"]
            self.assertEqual(len(listings), 2)
            self.assertEqual(listings[1]["params"]["cursor"], "page2")
            result = await client.call_tool(server, "get_order", {"order_id": "A-100"})

        self.assertEqual(tools[0]["name"], "get_order")
        self.assertEqual(result["text"], "已发货")
        self.assertEqual(result["structuredContent"], {"status": "shipped"})
        call = next(item for item in requests if item["method"] == "tools/call")
        self.assertEqual(call["params"]["arguments"], {"order_id": "A-100"})

    async def test_sdk_timeout_has_readable_error(self):
        @asynccontextmanager
        async def timed_out(*args):
            raise httpx2.ReadTimeout("read timed out")
            yield

        server = RuntimeMcpServer(id="mcp_timeout", url="https://mcp.example.com/mcp")
        client = McpToolClient()
        with patch.object(client, "_transport_streams", timed_out):
            with self.assertRaisesRegex(McpClientError, "超时"):
                await client.list_tools(server)

    def test_server_schema_accepts_sse(self):
        server = McpServerUpsert(
            name="旧版 MCP 服务",
            transport="sse",
            url="https://mcp.example.com/sse",
        )

        self.assertEqual(server.transport, "sse")

    async def test_sse_transport_uses_sdk_sse_client(self):
        captured: dict = {}

        @asynccontextmanager
        async def fake_sse_client(url, **kwargs):
            captured["url"] = url
            captured.update(kwargs)
            yield "read-stream", "write-stream"

        server = RuntimeMcpServer(
            id="mcp_sse",
            transport="sse",
            url="https://mcp.example.com/sse",
        )
        client = McpToolClient()

        with patch("services.tools.mcp_client.sse_client", fake_sse_client):
            async with client._transport_streams(
                server,
                {"Authorization": "Bearer test"},
                20,
            ) as streams:
                self.assertEqual(streams, ("read-stream", "write-stream"))

        self.assertEqual(captured["url"], server.url)
        self.assertEqual(captured["headers"], {"Authorization": "Bearer test"})
        self.assertEqual(captured["timeout"], 20)
        self.assertEqual(captured["sse_read_timeout"], 20)


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