refactor: update test files to use ParsedMessage

Updates test files to work with the ParsedMessage stream type aliases
and fixes a line length issue in test_201_client_hangs_on_logging.py.

Github-Issue:#201

Author: dsp-ant

tests/client/test_session.py

@@ -2,6 +2,7 @@
 import pytest
 
 from mcp.client.session import ClientSession
+from mcp.shared.session import ParsedMessage
 from mcp.types import (
     LATEST_PROTOCOL_VERSION,
     ClientNotification,
@@ -11,7 +12,6 @@
     InitializeRequest,
     InitializeResult,
     JSONRPCMessage,
-    JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
     ServerCapabilities,
@@ -22,10 +22,10 @@
 @pytest.mark.anyio
 async def test_client_session_initialize():
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        ParsedMessage[None]
     ](1)
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        ParsedMessage[None]
     ](1)
 
     initialized_notification = None
@@ -57,20 +57,23 @@ async def mock_server():
 
         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
-                        ),
-                    )
+                ParsedMessage(
+                    root=JSONRPCMessage(
+                        JSONRPCResponse(
+                            jsonrpc="2.0",
+                            id=jsonrpc_request.root.id,
+                            result=result.model_dump(
+                                by_alias=True, mode="json", exclude_none=True
+                            ),
+                        )
+                    ),
+                    raw=None,
                 )
             )
             jsonrpc_notification = await client_to_server_receive.receive()
-            assert isinstance(jsonrpc_notification.root, JSONRPCNotification)
+            assert isinstance(jsonrpc_notification.root, ParsedMessage)
             initialized_notification = ClientNotification.model_validate(
-                jsonrpc_notification.model_dump(
+                jsonrpc_notification.root.model_dump(
                     by_alias=True, mode="json", exclude_none=True
                 )
             )

tests/server/test_lifespan.py

@@ -117,7 +117,7 @@ async def run_server():
 
         # Get response and verify
         response = await receive_stream2.receive()
-        assert response.root.result["content"][0]["text"] == "true"
+        assert response.root.root.result["content"][0]["text"] == "true"
 
         # Cancel server task
         tg.cancel_scope.cancel()
@@ -213,7 +213,7 @@ async def run_server():
 
         # Get response and verify
         response = await receive_stream2.receive()
-        assert response.root.result["content"][0]["text"] == "true"
+        assert response.root.root.result["content"][0]["text"] == "true"
 
         # Cancel server task
         tg.cancel_scope.cancel()

tests/server/test_session.py

@@ -6,10 +6,10 @@
 from mcp.server.lowlevel import NotificationOptions
 from mcp.server.models import InitializationOptions
 from mcp.server.session import ServerSession
+from mcp.shared.session import ParsedMessage
 from mcp.types import (
     ClientNotification,
     InitializedNotification,
-    JSONRPCMessage,
     PromptsCapability,
     ResourcesCapability,
     ServerCapabilities,
@@ -19,10 +19,10 @@
 @pytest.mark.anyio
 async def test_server_session_initialize():
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        ParsedMessage[None]
     ](1)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        ParsedMessage[None]
     ](1)
 
     async def run_client(client: ClientSession):

tests/server/test_stdio.py

@@ -4,6 +4,7 @@
 import pytest
 
 from mcp.server.stdio import stdio_server
+from mcp.shared.session import ParsedMessage
 from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
 
 
@@ -13,8 +14,12 @@ async def test_stdio_server():
     stdout = io.StringIO()
 
     messages = [
-        JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")),
-        JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})),
+        ParsedMessage(
+            root=JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"))
+        ),
+        ParsedMessage(
+            root=JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={}))
+        ),
     ]
 
     for message in messages:
@@ -35,17 +40,25 @@ async def test_stdio_server():
 
         # 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[0] == ParsedMessage(
+            root=JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"))
         )
-        assert received_messages[1] == JSONRPCMessage(
-            root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})
+        assert received_messages[1] == ParsedMessage(
+            root=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={})),
+            ParsedMessage(
+                root=JSONRPCMessage(
+                    root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")
+                )
+            ),
+            ParsedMessage(
+                root=JSONRPCMessage(
+                    root=JSONRPCResponse(jsonrpc="2.0", id=4, result={})
+                )
+            ),
         ]
 
         async with write_stream:
@@ -57,12 +70,12 @@ async def test_stdio_server():
     assert len(output_lines) == 2
 
     received_responses = [
-        JSONRPCMessage.model_validate_json(line.strip()) for line in output_lines
+        ParsedMessage.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[0] == ParsedMessage(
+        root=JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping"))
     )
-    assert received_responses[1] == JSONRPCMessage(
-        root=JSONRPCResponse(jsonrpc="2.0", id=4, result={})
+    assert received_responses[1] == ParsedMessage(
+        root=JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=4, result={}))
     )