Making unit tests simpler

Author: aravind-segu

src/mcp/client/streamable_http.py

@@ -42,7 +42,6 @@
 MCP_SESSION_ID = "mcp-session-id"
 LAST_EVENT_ID = "last-event-id"
 CONTENT_TYPE = "content-type"
-HEADER_CAPTURE = "[TESTING_HEADER_CAPTURE]"
 ACCEPT = "Accept"
 
 
@@ -273,16 +272,6 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
                 if is_complete:
                     break
 
-    async def _is_testing_header_capture(self, response: httpx.Response) -> str | None:
-        try:
-            content = await response.aread()
-            if content.decode().startswith(HEADER_CAPTURE):
-                return content.decode()
-        except Exception as _:
-            return None
-
-        return None
-
     async def _handle_post_request(self, ctx: RequestContext) -> None:
         """Handle a POST request with response processing."""
         headers = await self._update_headers(ctx.headers)
@@ -307,23 +296,6 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
                     )
                 return
 
-            # To test if headers are being forwarded correctly, in unit tests
-            # we have a mock server that returns a 418 status code with the
-            # HEADER_CAPTURE prefix. If the response has this status code
-            # with the prefix, return the response content as part of the error message.
-            if response.status_code == 418:
-                test_error_message = await self._is_testing_header_capture(response)
-                # If this is coming from the test case return the response content
-                if test_error_message and isinstance(message.root, JSONRPCRequest):
-                    jsonrpc_error = JSONRPCError(
-                        jsonrpc="2.0",
-                        id=message.root.id,
-                        error=ErrorData(code=32600, message=test_error_message),
-                    )
-                    session_message = SessionMessage(JSONRPCMessage(jsonrpc_error))
-                    await ctx.read_stream_writer.send(session_message)
-                    return
-
             response.raise_for_status()
             if is_initialization:
                 self._maybe_extract_session_id_from_response(response)

tests/shared/test_streamable_http.py

@@ -4,7 +4,6 @@
 Contains tests for both server and client sides of the StreamableHTTP transport.
 """
 
-import json
 import multiprocessing
 import socket
 import time
@@ -19,13 +18,11 @@
 import uvicorn
 from pydantic import AnyUrl
 from starlette.applications import Starlette
-from starlette.requests import Request
-from starlette.responses import Response
 from starlette.routing import Mount
 
 import mcp.types as types
 from mcp.client.session import ClientSession
-from mcp.client.streamable_http import HEADER_CAPTURE, streamablehttp_client
+from mcp.client.streamable_http import streamablehttp_client
 from mcp.server import Server
 from mcp.server.streamable_http import (
     MCP_SESSION_ID_HEADER,
@@ -247,46 +244,8 @@ def create_app(
     return app
 
 
-def create_header_capture_app() -> Starlette:
-    """Implement a minimal Starlette app that intercepts every request,
-    extracts its headers, and responds with status 418 (Test Status code),
-    embedding the captured headers as the JSON response body.
-    We use this server solely to verify that the MCP Server is forwarding
-    headers correctly."""
-
-    # Create a wrapper that captures headers and returns them in error response
-    async def header_capture_wrapper(scope, receive, send):
-        # Capture headers
-        request = Request(scope, receive=receive)
-        headers = dict(request.headers)
-
-        # Return error response with headers in body
-        response = Response(
-            HEADER_CAPTURE + json.dumps({"headers": headers}),
-            status_code=418,
-        )
-        await response(scope, receive, send)
-
-    # Create an ASGI application that uses our wrapper
-    app = Starlette(
-        debug=True,
-        routes=[
-            Mount("/mcp", app=header_capture_wrapper),
-        ],
-    )
-
-    return app
-
-
-def _get_captured_headrs(str) -> dict[str, str]:
-    return json.loads(str.split(HEADER_CAPTURE)[1])["headers"]
-
-
 def run_server(
-    port: int,
-    is_json_response_enabled=False,
-    event_store: EventStore | None = None,
-    testing_header_capture: bool = False,
+    port: int, is_json_response_enabled=False, event_store: EventStore | None = None
 ) -> None:
     """Run the test server.
 
@@ -296,11 +255,7 @@ def run_server(
         event_store: Optional event store for testing resumability.
     """
 
-    if testing_header_capture:
-        app = create_header_capture_app()
-    else:
-        app = create_app(is_json_response_enabled, event_store)
-
+    app = create_app(is_json_response_enabled, event_store)
     # Configure server
     config = uvicorn.Config(
         app=app,
@@ -341,16 +296,11 @@ def json_server_port() -> int:
         return s.getsockname()[1]
 
 
-def _start_basic_server(
-    basic_server_port: int, testing_header_capture: bool
-) -> Generator[None, None, None]:
+@pytest.fixture
+def basic_server(basic_server_port: int) -> Generator[None, None, None]:
+    """Start a basic server."""
     proc = multiprocessing.Process(
-        target=run_server,
-        kwargs={
-            "port": basic_server_port,
-            "testing_header_capture": testing_header_capture,
-        },
-        daemon=True,
+        target=run_server, kwargs={"port": basic_server_port}, daemon=True
     )
     proc.start()
 
@@ -375,18 +325,6 @@ def _start_basic_server(
     proc.join(timeout=2)
 
 
-@pytest.fixture
-def basic_server(basic_server_port: int) -> Generator[None, None, None]:
-    yield from _start_basic_server(basic_server_port, testing_header_capture=False)
-
-
-@pytest.fixture
-def basic_server_with_header_capture(
-    basic_server_port: int,
-) -> Generator[None, None, None]:
-    yield from _start_basic_server(basic_server_port, testing_header_capture=True)
-
-
 @pytest.fixture
 def event_store() -> SimpleEventStore:
     """Create a test event store."""
@@ -1295,83 +1233,55 @@ def __init__(self, token: str):
         self.token = token
 
     async def get_auth_headers(self) -> dict[str, str]:
-        return {"Authorization": f"Bearer {self.token}"}
+        return {"Authorization": "Bearer " + self.token}
 
 
 @pytest.mark.anyio
-async def test_auth_client_provider_headers(
-    basic_server_with_header_capture, basic_server_url
-):
+async def test_auth_client_provider_headers(basic_server, basic_server_url):
     """Test that auth token provider correctly sets Authorization header."""
     # Create a mock token provider
-    client_provider = MockAuthClientProvider("short-lived-token-123")
+    client_provider = MockAuthClientProvider("test-token-123")
+    client_provider.get_auth_headers = AsyncMock(
+        return_value={"Authorization": "Bearer test-token-123"}
+    )
 
     # Create client with token provider
     async with streamablehttp_client(
         f"{basic_server_url}/mcp", auth_client_provider=client_provider
     ) as (read_stream, write_stream, _):
         async with ClientSession(read_stream, write_stream) as session:
             # Initialize the session
-            with pytest.raises(McpError) as mcpError:
-                _ = await session.initialize()
-                assert (
-                    _get_captured_headrs(mcpError.value.error.message)["Authorization"]
-                    == "Bearer short-lived-token-123"
-                )
+            result = await session.initialize()
+            assert isinstance(result, InitializeResult)
+
+            # Make a request to verify headers
+            tools = await session.list_tools()
+            assert len(tools.tools) == 4
+
+    client_provider.get_auth_headers.assert_called()
 
 
 @pytest.mark.anyio
-async def test_auth_client_provider_token_called_on_every_request(
-    basic_server_with_header_capture, basic_server_url
-):
+async def test_auth_client_provider_called_per_request(basic_server, basic_server_url):
     """Test that auth token provider can return different tokens."""
     # Create a dynamic token provider
-    client_provider = MockAuthClientProvider("short-lived-token-123")
+    client_provider = MockAuthClientProvider("test-token-123")
+    client_provider.get_auth_headers = AsyncMock(
+        return_value={"Authorization": "Bearer test-token-123"}
+    )
 
+    # Create client with dynamic token provider
     async with streamablehttp_client(
         f"{basic_server_url}/mcp", auth_client_provider=client_provider
     ) as (read_stream, write_stream, _):
         async with ClientSession(read_stream, write_stream) as session:
             # Initialize the session
-            with pytest.raises(McpError) as mcpError:
-                _ = await session.initialize()
-                assert (
-                    _get_captured_headrs(mcpError.value.error.message)["Authorization"]
-                    == "Bearer short-lived-token-123"
-                )
-
-            # Mock a new token and ensure the new token is returned
-            client_provider.get_auth_headers = AsyncMock(
-                return_value={"Authorization": "Bearer short-lived-token-456"}
-            )
-            with pytest.raises(McpError) as mcpError:
-                _ = await session.initialize()
-                assert (
-                    _get_captured_headrs(mcpError.value.error.message)["Authorization"]
-                    == "Bearer short-lived-token-456"
-                )
+            result = await session.initialize()
+            assert isinstance(result, InitializeResult)
 
+            # Make multiple requests to verify token updates
+            for i in range(3):
+                tools = await session.list_tools()
+                assert len(tools.tools) == 4
 
-@pytest.mark.anyio
-async def test_auth_client_provider_headers_not_overridden(
-    basic_server_with_header_capture, basic_server_url
-):
-    """Test that provided headers override auth client provider headers."""
-    # Create a mock token provider
-    client_provider = MockAuthClientProvider("short-lived-token")
-
-    # Create client with token provider and custom headers
-    custom_headers = {"Authorization": "Bearer original-long-lived-token"}
-    async with streamablehttp_client(
-        f"{basic_server_url}/mcp",
-        auth_client_provider=client_provider,
-        headers=custom_headers,
-    ) as (read_stream, write_stream, _):
-        async with ClientSession(read_stream, write_stream) as session:
-            # Original token is used and not short-lived-token from the provider
-            with pytest.raises(McpError) as mcpError:
-                _ = await session.initialize()
-                assert (
-                    _get_captured_headrs(mcpError.value.error.message)["Authorization"]
-                    == "Bearer original-long-lived-token"
-                )
+    client_provider.get_auth_headers.call_count > 1