Support different transports in Client (#1972)

tests/client/test_client.py

@@ -2,11 +2,14 @@
 
 from __future__ import annotations
 
+from unittest.mock import patch
+
 import anyio
 import pytest
 from inline_snapshot import snapshot
 
 import mcp.types as types
+from mcp.client._memory import InMemoryTransport
 from mcp.client.client import Client
 from mcp.server import Server
 from mcp.server.mcpserver import MCPServer
@@ -285,3 +288,21 @@ async def test_complete_with_prompt_reference(simple_server: Server):
         ref = types.PromptReference(type="ref/prompt", name="test_prompt")
         result = await client.complete(ref=ref, argument={"name": "arg", "value": "test"})
         assert result == snapshot(types.CompleteResult(completion=types.Completion(values=[])))
+
+
+def test_client_with_url_initializes_streamable_http_transport():
+    with patch("mcp.client.client.streamable_http_client") as mock:
+        _ = Client("http://localhost:8000/mcp")
+    mock.assert_called_once_with("http://localhost:8000/mcp")
+
+
+async def test_client_uses_transport_directly(app: MCPServer):
+    transport = InMemoryTransport(app)
+    async with Client(transport) as client:
+        result = await client.call_tool("greet", {"name": "Transport"})
+        assert result == snapshot(
+            CallToolResult(
+                content=[TextContent(text="Hello, Transport!")],
+                structured_content={"result": "Hello, Transport!"},
+            )
+        )

tests/client/test_http_unicode.py

@@ -173,7 +173,7 @@ async def test_streamable_http_client_unicode_tool_call(running_unicode_server:
     base_url = running_unicode_server
     endpoint_url = f"{base_url}/mcp"
 
-    async with streamable_http_client(endpoint_url) as (read_stream, write_stream, _get_session_id):
+    async with streamable_http_client(endpoint_url) as (read_stream, write_stream):
         async with ClientSession(read_stream, write_stream) as session:
             await session.initialize()
 
@@ -205,7 +205,7 @@ async def test_streamable_http_client_unicode_prompts(running_unicode_server: st
     base_url = running_unicode_server
     endpoint_url = f"{base_url}/mcp"
 
-    async with streamable_http_client(endpoint_url) as (read_stream, write_stream, _get_session_id):
+    async with streamable_http_client(endpoint_url) as (read_stream, write_stream):
         async with ClientSession(read_stream, write_stream) as session:
             await session.initialize()
 

tests/client/test_notification_response.py

@@ -125,12 +125,8 @@ async def test_non_compliant_notification_response(non_sdk_server: None, non_sdk
         if isinstance(message, Exception):
             returned_exception = message
 
-    async with streamable_http_client(server_url) as (read_stream, write_stream, _):
-        async with ClientSession(
-            read_stream,
-            write_stream,
-            message_handler=message_handler,
-        ) as session:
+    async with streamable_http_client(server_url) as (read_stream, write_stream):
+        async with ClientSession(read_stream, write_stream, message_handler=message_handler) as session:
             # Initialize should work normally
             await session.initialize()
 

tests/client/test_session_group.py

@@ -278,6 +278,7 @@ async def test_client_session_group_disconnect_non_existent_server():
         await group.disconnect_from_server(session)
 
 
+# TODO(Marcelo): This is horrible. We should drop this test.
 @pytest.mark.anyio
 @pytest.mark.parametrize(
     "server_params_instance, client_type_name, patch_target_for_client_func",
@@ -310,19 +311,8 @@ async def test_client_session_group_establish_session_parameterized(
             mock_read_stream = mock.AsyncMock(name=f"{client_type_name}Read")
             mock_write_stream = mock.AsyncMock(name=f"{client_type_name}Write")
 
-            # streamable_http_client's __aenter__ returns three values
-            if client_type_name == "streamablehttp":
-                mock_extra_stream_val = mock.AsyncMock(name="StreamableExtra")
-                mock_client_cm_instance.__aenter__.return_value = (
-                    mock_read_stream,
-                    mock_write_stream,
-                    mock_extra_stream_val,
-                )
-            else:
-                mock_client_cm_instance.__aenter__.return_value = (
-                    mock_read_stream,
-                    mock_write_stream,
-                )
+            # All client context managers return (read_stream, write_stream)
+            mock_client_cm_instance.__aenter__.return_value = (mock_read_stream, mock_write_stream)
 
             mock_client_cm_instance.__aexit__ = mock.AsyncMock(return_value=None)
             mock_specific_client_func.return_value = mock_client_cm_instance

tests/client/transports/test_memory.py

@@ -53,7 +53,7 @@ pytestmark = pytest.mark.anyio
 async def test_with_server(simple_server: Server):
     """Test creating transport with a Server instance."""
     transport = InMemoryTransport(simple_server)
-    async with transport.connect() as (read_stream, write_stream):
+    async with transport as (read_stream, write_stream):
         assert read_stream is not None
         assert write_stream is not None
 
@@ -61,7 +61,7 @@ async def test_with_server(simple_server: Server):
 async def test_with_mcpserver(mcpserver_server: MCPServer):
     """Test creating transport with an MCPServer instance."""
     transport = InMemoryTransport(mcpserver_server)
-    async with transport.connect() as (read_stream, write_stream):
+    async with transport as (read_stream, write_stream):
         assert read_stream is not None
         assert write_stream is not None
 
@@ -93,5 +93,5 @@ async def test_call_tool(mcpserver_server: MCPServer):
 async def test_raise_exceptions(mcpserver_server: MCPServer):
     """Test that raise_exceptions parameter is passed through."""
     transport = InMemoryTransport(mcpserver_server, raise_exceptions=True)
-    async with transport.connect() as (read_stream, _write_stream):
+    async with transport as (read_stream, _write_stream):
         assert read_stream is not None