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/client/__init__.py Normal file
View File

View File

@@ -0,0 +1,93 @@
import anyio
import pytest
from mcp_python.client.session import ClientSession
from mcp_python.types import (
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")
)
assert isinstance(request.root, InitializeRequest)
result = ServerResult(
InitializeResult(
protocolVersion=1,
capabilities=ServerCapabilities(
logging=None,
resources=None,
tools=None,
experimental=None,
prompts=None,
),
serverInfo=Implementation(name="mock-server", version="0.1.0"),
)
)
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"),
)
)
)
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")
)
async def listen_session():
async for message in session.incoming_messages:
if isinstance(message, Exception):
raise message
async with (
ClientSession(server_to_client_receive, client_to_server_send) as session,
anyio.create_task_group() as tg,
):
tg.start_soon(mock_server)
tg.start_soon(listen_session)
result = await session.initialize()
# Assert the result
assert isinstance(result, InitializeResult)
assert result.protocolVersion == 1
assert isinstance(result.capabilities, ServerCapabilities)
assert result.serverInfo == Implementation(name="mock-server", version="0.1.0")
# Check that the client sent the initialized notification
assert initialized_notification
assert isinstance(initialized_notification.root, InitializedNotification)

View File

@@ -0,0 +1,38 @@
import pytest
from mcp_python.client.stdio import StdioServerParameters, stdio_client
from mcp_python.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
@pytest.mark.anyio
async def test_stdio_client():
server_parameters = StdioServerParameters(command="/usr/bin/tee")
async with stdio_client(server_parameters) as (read_stream, write_stream):
# Test sending and receiving messages
messages = [
JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")),
JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})),
]
async with write_stream:
for message in messages:
await write_stream.send(message)
read_messages = []
async with read_stream:
async for message in read_stream:
if isinstance(message, Exception):
raise message
read_messages.append(message)
if len(read_messages) == 2:
break
assert len(read_messages) == 2
assert read_messages[0] == JSONRPCMessage(
root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
)
assert read_messages[1] == JSONRPCMessage(
root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})
)