Initial import

This commit is contained in:
David Soria Parra
2024-09-24 22:04:19 +01:00
commit 4cbf815430
33 changed files with 2986 additions and 0 deletions

0
tests/server/__init__.py Normal file
View File

View File

@@ -0,0 +1,59 @@
import anyio
import pytest
from mcp_python.client.session import ClientSession
from mcp_python.server.session import ServerSession
from mcp_python.types import (
ClientNotification,
InitializedNotification,
JSONRPCMessage,
)
@pytest.mark.anyio
async def test_server_session_initialize():
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream(
1, JSONRPCMessage
)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream(
1, JSONRPCMessage
)
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
) 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

View File

@@ -0,0 +1,68 @@
import io
import anyio
import pytest
from mcp_python.server.stdio import stdio_server
from mcp_python.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() + "\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={})
)