forked from modelcontextprotocol/python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_stdio.py
More file actions
68 lines (56 loc) · 2.18 KB
/
test_stdio.py
File metadata and controls
68 lines (56 loc) · 2.18 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
import io
import anyio
import pytest
from mcp.server.stdio import stdio_server
from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
@pytest.mark.anyio
async def test_stdio_server():
stdin = io.StringIO()
stdout = io.StringIO()
messages = [
JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")),
JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})),
]
for message in messages:
stdin.write(message.model_dump_json(by_alias=True, exclude_none=True) + "\n")
stdin.seek(0)
async with stdio_server(
stdin=anyio.AsyncFile(stdin), stdout=anyio.AsyncFile(stdout)
) as (read_stream, write_stream):
received_messages = []
async with read_stream:
async for message in read_stream:
if isinstance(message, Exception):
raise message
received_messages.append(message)
if len(received_messages) == 2:
break
# Verify received messages
assert len(received_messages) == 2
assert received_messages[0] == JSONRPCMessage(
root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
)
assert received_messages[1] == JSONRPCMessage(
root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})
)
# Test sending responses from the server
responses = [
JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")),
JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=4, result={})),
]
async with write_stream:
for response in responses:
await write_stream.send(response)
stdout.seek(0)
output_lines = stdout.readlines()
assert len(output_lines) == 2
received_responses = [
JSONRPCMessage.model_validate_json(line.strip()) for line in output_lines
]
assert len(received_responses) == 2
assert received_responses[0] == JSONRPCMessage(
root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")
)
assert received_responses[1] == JSONRPCMessage(
root=JSONRPCResponse(jsonrpc="2.0", id=4, result={})
)