128 lines
4.8 KiB
Python
128 lines
4.8 KiB
Python
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()
|