forked from modelcontextprotocol/python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_session.py
More file actions
113 lines (101 loc) · 3.74 KB
/
Copy pathtest_session.py
File metadata and controls
113 lines (101 loc) · 3.74 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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import anyio
import pytest
import mcp.types as types
from mcp.client.session import ClientSession
from mcp.shared.session import RequestResponder
from mcp.types import (
LATEST_PROTOCOL_VERSION,
ClientNotification,
ClientRequest,
Implementation,
InitializedNotification,
InitializeRequest,
InitializeResult,
JSONRPCMessage,
JSONRPCNotification,
JSONRPCRequest,
JSONRPCResponse,
ServerCapabilities,
ServerResult,
)
@pytest.mark.anyio
async def test_client_session_initialize():
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
JSONRPCMessage
](1)
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
JSONRPCMessage
](1)
initialized_notification = None
async def mock_server():
nonlocal initialized_notification
jsonrpc_request = await client_to_server_receive.receive()
assert isinstance(jsonrpc_request.root, JSONRPCRequest)
request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
)
assert isinstance(request.root, InitializeRequest)
result = ServerResult(
InitializeResult(
protocolVersion=LATEST_PROTOCOL_VERSION,
capabilities=ServerCapabilities(
logging=None,
resources=None,
tools=None,
experimental=None,
prompts=None,
),
serverInfo=Implementation(name="mock-server", version="0.1.0"),
instructions="The server instructions.",
)
)
async with server_to_client_send:
await server_to_client_send.send(
JSONRPCMessage(
JSONRPCResponse(
jsonrpc="2.0",
id=jsonrpc_request.root.id,
result=result.model_dump(
by_alias=True, mode="json", exclude_none=True
),
)
)
)
jsonrpc_notification = await client_to_server_receive.receive()
assert isinstance(jsonrpc_notification.root, JSONRPCNotification)
initialized_notification = ClientNotification.model_validate(
jsonrpc_notification.model_dump(
by_alias=True, mode="json", exclude_none=True
)
)
# Create a message handler to catch exceptions
async def message_handler(
message: RequestResponder[types.ServerRequest, types.ClientResult]
| types.ServerNotification
| Exception,
) -> None:
if isinstance(message, Exception):
raise message
async with (
ClientSession(
server_to_client_receive,
client_to_server_send,
message_handler=message_handler,
) as session,
anyio.create_task_group() as tg,
client_to_server_send,
client_to_server_receive,
server_to_client_send,
server_to_client_receive,
):
tg.start_soon(mock_server)
result = await session.initialize()
# Assert the result
assert isinstance(result, InitializeResult)
assert result.protocolVersion == LATEST_PROTOCOL_VERSION
assert isinstance(result.capabilities, ServerCapabilities)
assert result.serverInfo == Implementation(name="mock-server", version="0.1.0")
assert result.instructions == "The server instructions."
# Check that the client sent the initialized notification
assert initialized_notification
assert isinstance(initialized_notification.root, InitializedNotification)