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
101 lines (83 loc) · 3.06 KB
/
test_session.py
File metadata and controls
101 lines (83 loc) · 3.06 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
import anyio
import pytest
from mcp.client.session import ClientSession
from mcp.server import Server
from mcp.server.lowlevel import NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp.types import (
ClientNotification,
InitializedNotification,
JSONRPCMessage,
PromptsCapability,
ResourcesCapability,
ServerCapabilities,
)
@pytest.mark.anyio
async def test_server_session_initialize():
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
JSONRPCMessage
](1)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
JSONRPCMessage
](1)
async def run_client(client: ClientSession):
async for message in client_session.incoming_messages:
if isinstance(message, Exception):
raise message
received_initialized = False
async def run_server():
nonlocal received_initialized
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="mcp",
server_version="0.1.0",
capabilities=ServerCapabilities(),
),
) as server_session:
async for message in server_session.incoming_messages:
if isinstance(message, Exception):
raise message
if isinstance(message, ClientNotification) and isinstance(
message.root, InitializedNotification
):
received_initialized = True
return
try:
async with (
ClientSession(
server_to_client_receive, client_to_server_send
) as client_session,
anyio.create_task_group() as tg,
):
tg.start_soon(run_client, client_session)
tg.start_soon(run_server)
await client_session.initialize()
except anyio.ClosedResourceError:
pass
assert received_initialized
@pytest.mark.anyio
async def test_server_capabilities():
server = Server("test")
notification_options = NotificationOptions()
experimental_capabilities = {}
# Initially no capabilities
caps = server.get_capabilities(notification_options, experimental_capabilities)
assert caps.prompts is None
assert caps.resources is None
# Add a prompts handler
@server.list_prompts()
async def list_prompts():
return []
caps = server.get_capabilities(notification_options, experimental_capabilities)
assert caps.prompts == PromptsCapability(listChanged=False)
assert caps.resources is None
# Add a resources handler
@server.list_resources()
async def list_resources():
return []
caps = server.get_capabilities(notification_options, experimental_capabilities)
assert caps.prompts == PromptsCapability(listChanged=False)
assert caps.resources == ResourcesCapability(subscribe=False, listChanged=False)