Initial import

tests/client/test_session.py

@@ -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)

tests/client/test_stdio.py

@@ -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={})
+        )

tests/server/test_session.py

@@ -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

tests/server/test_stdio.py

@@ -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={})
+    )

tests/test_types.py

@@ -0,0 +1,24 @@
+from mcp_python.types import ClientRequest, JSONRPCMessage, JSONRPCRequest
+
+
+def test_jsonrpc_request():
+    json_data = {
+        "jsonrpc": "2.0",
+        "id": 1,
+        "method": "initialize",
+        "params": {
+            "protocolVersion": 1,
+            "capabilities": {"batch": None, "sampling": None},
+            "clientInfo": {"name": "mcp_python", "version": "0.1.0"},
+        },
+    }
+
+    request = JSONRPCMessage.model_validate(json_data)
+    assert isinstance(request.root, JSONRPCRequest)
+    ClientRequest.model_validate(request.model_dump(by_alias=True))
+
+    assert request.root.jsonrpc == "2.0"
+    assert request.root.id == 1
+    assert request.root.method == "initialize"
+    assert request.root.params is not None
+    assert request.root.params["protocolVersion"] == 1