wehub-resource-sync
231 lines
6.8 KiB
Python
231 lines
6.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from opensquilla.mcp.stdio import MCPStdioClient
|
|
from opensquilla.mcp.types import MCPServerConfig
|
|
|
|
_SDK_SERVER_SCRIPT = str(Path(__file__).parent / "fixtures" / "fastmcp_server.py")
|
|
|
|
# A stdio MCP server (newline-delimited JSON-RPC) exposing one "search" tool
|
|
# whose call response carries the MCP result-level ``isError`` flag.
|
|
_ERROR_SERVER_SCRIPT = r"""
|
|
import json, sys
|
|
|
|
def send(payload):
|
|
sys.stdout.buffer.write((json.dumps(payload) + "\n").encode())
|
|
sys.stdout.buffer.flush()
|
|
|
|
while True:
|
|
line = sys.stdin.buffer.readline()
|
|
if not line:
|
|
break
|
|
msg = json.loads(line.decode())
|
|
msg_id = msg.get("id")
|
|
if msg_id is None:
|
|
continue # notification
|
|
method = msg.get("method")
|
|
if method == "initialize":
|
|
send({"jsonrpc": "2.0", "id": msg_id,
|
|
"result": {"protocolVersion": "2024-11-05", "capabilities": {},
|
|
"serverInfo": {"name": "fake", "version": "0.0.1"}}})
|
|
elif method == "tools/call":
|
|
send({"jsonrpc": "2.0", "id": msg_id,
|
|
"result": {"isError": True,
|
|
"content": [{"type": "text",
|
|
"text": "Error: upstream API rejected the query"}]}})
|
|
else:
|
|
send({"jsonrpc": "2.0", "id": msg_id, "result": {}})
|
|
"""
|
|
|
|
|
|
class _FakeProcess:
|
|
def __init__(self, *, exits_on_terminate: bool = True) -> None:
|
|
self.returncode: int | None = None
|
|
self.terminated = False
|
|
self.killed = False
|
|
self.wait_calls = 0
|
|
self.exits_on_terminate = exits_on_terminate
|
|
|
|
def terminate(self) -> None:
|
|
self.terminated = True
|
|
if self.exits_on_terminate:
|
|
self.returncode = 0
|
|
|
|
def kill(self) -> None:
|
|
self.killed = True
|
|
self.returncode = -9
|
|
|
|
async def wait(self) -> int:
|
|
self.wait_calls += 1
|
|
if self.returncode is None:
|
|
await asyncio.sleep(3600)
|
|
return self.returncode
|
|
|
|
|
|
def _client_with_process(process: _FakeProcess) -> MCPStdioClient:
|
|
client = MCPStdioClient(MCPServerConfig(name="demo", transport="stdio", command="demo"))
|
|
client._process = process # type: ignore[assignment]
|
|
return client
|
|
|
|
|
|
class _RecordingStdin:
|
|
def __init__(self) -> None:
|
|
self.writes: list[bytes] = []
|
|
self.changed = asyncio.Event()
|
|
|
|
def write(self, data: bytes) -> None:
|
|
self.writes.append(data)
|
|
self.changed.set()
|
|
|
|
async def drain(self) -> None:
|
|
return None
|
|
|
|
async def wait_for_count(self, count: int) -> None:
|
|
while len(self.writes) < count:
|
|
self.changed.clear()
|
|
await self.changed.wait()
|
|
|
|
|
|
class _QueuedStdout:
|
|
def __init__(self) -> None:
|
|
self.lines: asyncio.Queue[bytes] = asyncio.Queue()
|
|
|
|
async def readline(self) -> bytes:
|
|
return await self.lines.get()
|
|
|
|
def respond(self, request_id: int) -> None:
|
|
self.lines.put_nowait(
|
|
(json.dumps({"jsonrpc": "2.0", "id": request_id, "result": {}}) + "\n").encode()
|
|
)
|
|
|
|
|
|
class _ConcurrentProcess:
|
|
def __init__(self) -> None:
|
|
self.stdin = _RecordingStdin()
|
|
self.stdout = _QueuedStdout()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_request_before_connect_reports_connection_error() -> None:
|
|
client = MCPStdioClient(MCPServerConfig(name="demo", transport="stdio", command="demo"))
|
|
|
|
with pytest.raises(ConnectionError, match="not connected"):
|
|
await client._send_request("tools/list")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_response_skips_invalid_utf8_lines() -> None:
|
|
process = _ConcurrentProcess()
|
|
process.stdout.lines.put_nowait(b"\xff\n")
|
|
process.stdout.respond(1)
|
|
client = _client_with_process(process) # type: ignore[arg-type]
|
|
|
|
assert await client._read_response(1) == {"jsonrpc": "2.0", "id": 1, "result": {}}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_requests_are_serialized_to_preserve_responses() -> None:
|
|
process = _ConcurrentProcess()
|
|
client = _client_with_process(process) # type: ignore[arg-type]
|
|
|
|
first = asyncio.create_task(client._send_request("first"))
|
|
second: asyncio.Task[dict] | None = None
|
|
try:
|
|
await process.stdin.wait_for_count(1)
|
|
second = asyncio.create_task(client._send_request("second"))
|
|
await asyncio.sleep(0)
|
|
|
|
assert len(process.stdin.writes) == 1
|
|
|
|
process.stdout.respond(1)
|
|
await first
|
|
await process.stdin.wait_for_count(2)
|
|
process.stdout.respond(2)
|
|
await second
|
|
finally:
|
|
first.cancel()
|
|
if second is not None:
|
|
second.cancel()
|
|
pending = [first] if second is None else [first, second]
|
|
await asyncio.gather(*pending, return_exceptions=True)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_waits_for_terminated_stdio_process() -> None:
|
|
process = _FakeProcess(exits_on_terminate=True)
|
|
|
|
await _client_with_process(process).close()
|
|
|
|
assert process.terminated is True
|
|
assert process.killed is False
|
|
assert process.wait_calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_kills_stdio_process_when_terminate_times_out(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
process = _FakeProcess(exits_on_terminate=False)
|
|
client = _client_with_process(process)
|
|
monkeypatch.setattr(client, "_CLOSE_TIMEOUT_SECONDS", 0.001)
|
|
|
|
await client.close()
|
|
|
|
assert process.terminated is True
|
|
assert process.killed is True
|
|
assert process.wait_calls == 2
|
|
|
|
|
|
def test_encode_request_is_newline_delimited_json() -> None:
|
|
encoded = MCPStdioClient._encode_request({"jsonrpc": "2.0", "id": 1, "method": "x"})
|
|
|
|
assert not encoded.startswith(b"Content-Length: ")
|
|
assert encoded.endswith(b"\n")
|
|
assert b"\n" not in encoded[:-1]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_connect_and_list_tools_against_sdk_stdio_server() -> None:
|
|
pytest.importorskip("mcp")
|
|
client = MCPStdioClient(
|
|
MCPServerConfig(
|
|
name="demo",
|
|
transport="stdio",
|
|
command=sys.executable,
|
|
args=[_SDK_SERVER_SCRIPT],
|
|
)
|
|
)
|
|
try:
|
|
await asyncio.wait_for(client.connect(), timeout=30.0)
|
|
tools = await asyncio.wait_for(client.list_tools(), timeout=30.0)
|
|
finally:
|
|
await client.close()
|
|
|
|
assert [t.name for t in tools] == ["ping"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_call_tool_honors_result_level_is_error_flag() -> None:
|
|
client = MCPStdioClient(
|
|
MCPServerConfig(
|
|
name="demo",
|
|
transport="stdio",
|
|
command=sys.executable,
|
|
args=["-c", _ERROR_SERVER_SCRIPT],
|
|
)
|
|
)
|
|
try:
|
|
await client.connect()
|
|
result = await client.call_tool("search", {"q": "x"})
|
|
finally:
|
|
await client.close()
|
|
|
|
assert "upstream API rejected" in result.content
|
|
assert result.is_error is True
|