refactor: use Client class in tests (#1900)

tests/client/test_list_methods_cursor.py

@@ -3,8 +3,7 @@ from collections.abc import Callable
 import pytest
 
 import mcp.types as types
-from mcp.client._memory import InMemoryTransport
+from mcp import Client
-from mcp.client.session import ClientSession
 from mcp.server import Server
 from mcp.server.fastmcp import FastMCP
 from mcp.types import ListToolsRequest, ListToolsResult
@@ -66,14 +65,11 @@ async def test_list_methods_params_parameter(
 
     See: https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/pagination#request-format
     """
-    transport = InMemoryTransport(full_featured_server)
+    async with Client(full_featured_server) as client:
-    async with transport.connect() as (read_stream, write_stream):
-        async with ClientSession(read_stream, write_stream) as session:
-            await session.initialize()
         spies = stream_spy()
 
         # Test without params (omitted)
-            method = getattr(session, method_name)
+        method = getattr(client, method_name)
         _ = await method()
         requests = spies.get_client_requests(method=request_method)
         assert len(requests) == 1
@@ -102,11 +98,8 @@ async def test_list_tools_with_strict_server_validation(
     full_featured_server: FastMCP,
 ):
     """Test pagination with a server that validates request format strictly."""
-    transport = InMemoryTransport(full_featured_server)
+    async with Client(full_featured_server) as client:
-    async with transport.connect() as (read_stream, write_stream):
+        result = await client.list_tools(params=types.PaginatedRequestParams())
-        async with ClientSession(read_stream, write_stream) as session:
-            await session.initialize()
-            result = await session.list_tools(params=types.PaginatedRequestParams())
         assert isinstance(result, ListToolsResult)
         assert len(result.tools) > 0
 
@@ -129,13 +122,9 @@ async def test_list_tools_with_lowlevel_server():
             ]
         )
 
-    transport = InMemoryTransport(server)
+    async with Client(server) as client:
-    async with transport.connect() as (read_stream, write_stream):
+        result = await client.list_tools(params=types.PaginatedRequestParams())
-        async with ClientSession(read_stream, write_stream) as session:
-            await session.initialize()
-
-            result = await session.list_tools(params=types.PaginatedRequestParams())
         assert result.tools[0].description == "cursor=None"
 
-            result = await session.list_tools(params=types.PaginatedRequestParams(cursor="page2"))
+        result = await client.list_tools(params=types.PaginatedRequestParams(cursor="page2"))
         assert result.tools[0].description == "cursor=page2"

tests/client/transports/test_memory.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from mcp import Client
 from mcp.client._memory import InMemoryTransport
 from mcp.server import Server
 from mcp.server.fastmcp import FastMCP
@@ -69,25 +70,14 @@ async def test_with_fastmcp(fastmcp_server: FastMCP):
 
 async def test_server_is_running(fastmcp_server: FastMCP):
     """Test that the server is running and responding to requests."""
-    from mcp.client.session import ClientSession
+    async with Client(fastmcp_server) as client:
-
+        assert client.server_capabilities is not None
-    transport = InMemoryTransport(fastmcp_server)
-    async with transport.connect() as (read_stream, write_stream):
-        async with ClientSession(read_stream, write_stream) as session:
-            result = await session.initialize()
-            assert result is not None
-            assert result.server_info.name == "test"
 
 
 async def test_list_tools(fastmcp_server: FastMCP):
     """Test listing tools through the transport."""
-    from mcp.client.session import ClientSession
+    async with Client(fastmcp_server) as client:
-
+        tools_result = await client.list_tools()
-    transport = InMemoryTransport(fastmcp_server)
-    async with transport.connect() as (read_stream, write_stream):
-        async with ClientSession(read_stream, write_stream) as session:
-            await session.initialize()
-            tools_result = await session.list_tools()
         assert len(tools_result.tools) > 0
         tool_names = [t.name for t in tools_result.tools]
         assert "greet" in tool_names
@@ -95,13 +85,8 @@ async def test_list_tools(fastmcp_server: FastMCP):
 
 async def test_call_tool(fastmcp_server: FastMCP):
     """Test calling a tool through the transport."""
-    from mcp.client.session import ClientSession
+    async with Client(fastmcp_server) as client:
-
+        result = await client.call_tool("greet", {"name": "World"})
-    transport = InMemoryTransport(fastmcp_server)
-    async with transport.connect() as (read_stream, write_stream):
-        async with ClientSession(read_stream, write_stream) as session:
-            await session.initialize()
-            result = await session.call_tool("greet", {"name": "World"})
         assert result is not None
         assert len(result.content) > 0
         assert "Hello, World!" in str(result.content[0])

tests/server/test_cancel_handling.py

@@ -6,8 +6,7 @@ import anyio
 import pytest
 
 import mcp.types as types
-from mcp.client._memory import InMemoryTransport
+from mcp import Client
-from mcp.client.session import ClientSession
 from mcp.server.lowlevel.server import Server
 from mcp.shared.exceptions import McpError
 from mcp.types import (
@@ -55,15 +54,11 @@ async def test_server_remains_functional_after_cancel():
             return [types.TextContent(type="text", text=f"Call number: {call_count}")]
         raise ValueError(f"Unknown tool: {name}")  # pragma: no cover
 
-    transport = InMemoryTransport(server)
+    async with Client(server) as client:
-    async with transport.connect() as (read_stream, write_stream):
-        async with ClientSession(read_stream, write_stream) as client:
-            await client.initialize()
-
         # First request (will be cancelled)
         async def first_request():
             try:
-                    await client.send_request(
+                await client.session.send_request(
                     ClientRequest(
                         CallToolRequest(
                             params=CallToolRequestParams(name="test_tool", arguments={}),
@@ -84,7 +79,7 @@ async def test_server_remains_functional_after_cancel():
 
             # Cancel it
             assert first_request_id is not None
-                await client.send_notification(
+            await client.session.send_notification(
                 ClientNotification(
                     CancelledNotification(
                         params=CancelledNotificationParams(
@@ -96,14 +91,7 @@ async def test_server_remains_functional_after_cancel():
             )
 
         # Second request (should work normally)
-            result = await client.send_request(
+        result = await client.call_tool("test_tool", {})
-                ClientRequest(
-                    CallToolRequest(
-                        params=CallToolRequestParams(name="test_tool", arguments={}),
-                    )
-                ),
-                CallToolResult,
-            )
 
         # Verify second request completed successfully
         assert len(result.content) == 1

tests/shared/test_progress_notifications.py

@@ -5,7 +5,7 @@ import anyio
 import pytest
 
 import mcp.types as types
-from mcp.client._memory import InMemoryTransport
+from mcp import Client
 from mcp.client.session import ClientSession
 from mcp.server import Server
 from mcp.server.lowlevel import NotificationOptions
@@ -369,21 +369,11 @@ async def test_progress_callback_exception_logging():
 
     # Test with mocked logging
     with patch("mcp.shared.session.logging.error", side_effect=mock_log_error):
-        transport = InMemoryTransport(server)
+        async with Client(server) as client:
-        async with transport.connect() as (read_stream, write_stream):
+            # Call tool with a failing progress callback
-            async with ClientSession(  # pragma: no branch
+            result = await client.call_tool(
-                read_stream=read_stream, write_stream=write_stream
+                "progress_tool",
-            ) as session:
+                arguments={},
-                await session.initialize()
-                # Send a request with a failing progress callback
-                result = await session.send_request(
-                    types.ClientRequest(
-                        types.CallToolRequest(
-                            method="tools/call",
-                            params=types.CallToolRequestParams(name="progress_tool", arguments={}),
-                        )
-                    ),
-                    types.CallToolResult,
                 progress_callback=failing_progress_callback,
             )
 

tests/shared/test_session.py

@@ -4,7 +4,7 @@ import anyio
 import pytest
 
 import mcp.types as types
-from mcp.client._memory import InMemoryTransport
+from mcp import Client
 from mcp.client.session import ClientSession
 from mcp.server.lowlevel.server import Server
 from mcp.shared.exceptions import McpError
@@ -25,38 +25,27 @@ from mcp.types import (
 )
 
 
-@pytest.fixture
-def mcp_server() -> Server:
-    return Server(name="test server")
-
-
 @pytest.mark.anyio
-async def test_in_flight_requests_cleared_after_completion(mcp_server: Server):
+async def test_in_flight_requests_cleared_after_completion():
     """Verify that _in_flight is empty after all requests complete."""
-    transport = InMemoryTransport(mcp_server)
+    server = Server(name="test server")
-    async with transport.connect() as (read_stream, write_stream):
+    async with Client(server) as client:
-        async with ClientSession(read_stream=read_stream, write_stream=write_stream) as session:
-            await session.initialize()
-
         # Send a request and wait for response
-            response = await session.send_ping()
+        response = await client.send_ping()
         assert isinstance(response, EmptyResult)
 
         # Verify _in_flight is empty
-            assert len(session._in_flight) == 0
+        assert len(client.session._in_flight) == 0
 
 
 @pytest.mark.anyio
 async def test_request_cancellation():
     """Test that requests can be cancelled while in-flight."""
-    # The tool is already registered in the fixture
-
     ev_tool_called = anyio.Event()
     ev_cancelled = anyio.Event()
     request_id = None
 
-    # Start the request in a separate task so we can cancel it
+    # Create a server with a slow tool
-    def make_server() -> Server:
     server = Server(name="TestSessionServer")
 
     # Register the tool handler
@@ -81,12 +70,10 @@ async def test_request_cancellation():
             )
         ]
 
-        return server
+    async def make_request(client: Client):
-
-    async def make_request(session: ClientSession):
         nonlocal ev_cancelled
         try:
-            await session.send_request(
+            await client.session.send_request(
                 ClientRequest(
                     types.CallToolRequest(
                         params=types.CallToolRequestParams(name="slow_tool", arguments={}),
@@ -100,12 +87,9 @@ async def test_request_cancellation():
             assert "Request cancelled" in str(e)
             ev_cancelled.set()
 
-    transport = InMemoryTransport(make_server())
+    async with Client(server) as client:
-    async with transport.connect() as (read_stream, write_stream):
-        async with ClientSession(read_stream=read_stream, write_stream=write_stream) as session:
-            await session.initialize()
         async with anyio.create_task_group() as tg:  # pragma: no branch
-                tg.start_soon(make_request, session)
+            tg.start_soon(make_request, client)
 
             # Wait for the request to be in-flight
             with anyio.fail_after(1):  # Timeout after 1 second
@@ -113,7 +97,7 @@ async def test_request_cancellation():
 
             # Send cancellation notification
             assert request_id is not None
-                await session.send_notification(
+            await client.session.send_notification(
                 ClientNotification(
                     CancelledNotification(
                         params=CancelledNotificationParams(request_id=request_id),