test: fix WS test port race; narrow to single smoke test covering both transport ends (#2267)

maxisbey · web-flow · commit 67201a9bbd9c · 2026-03-18T15:48:30.000Z
diff --git a/src/mcp/server/websocket.py b/src/mcp/server/websocket.py
@@ -10,7 +10,7 @@
 from mcp.shared.message import SessionMessage
 
 
-@asynccontextmanager  # pragma: no cover
+@asynccontextmanager
 async def websocket_server(scope: Scope, receive: Receive, send: Send):
     """WebSocket server transport for MCP. This is an ASGI application, suitable for use
     with a framework like Starlette and a server like Hypercorn.
@@ -34,13 +34,13 @@ async def ws_reader():
                 async for msg in websocket.iter_text():
                     try:
                         client_message = types.jsonrpc_message_adapter.validate_json(msg, by_name=False)
-                    except ValidationError as exc:
+                    except ValidationError as exc:  # pragma: no cover
                         await read_stream_writer.send(exc)
                         continue
 
                     session_message = SessionMessage(client_message)
                     await read_stream_writer.send(session_message)
-        except anyio.ClosedResourceError:
+        except anyio.ClosedResourceError:  # pragma: no cover
             await websocket.close()
 
     async def ws_writer():
@@ -49,7 +49,7 @@ async def ws_writer():
                 async for session_message in write_stream_reader:
                     obj = session_message.message.model_dump_json(by_alias=True, exclude_unset=True)
                     await websocket.send_text(obj)
-        except anyio.ClosedResourceError:
+        except anyio.ClosedResourceError:  # pragma: no cover
             await websocket.close()
 
     async with anyio.create_task_group() as tg:
diff --git a/tests/client/test_client.py b/tests/client/test_client.py
@@ -8,7 +8,7 @@
 import pytest
 from inline_snapshot import snapshot
 
-from mcp import types
+from mcp import MCPError, types
 from mcp.client._memory import InMemoryTransport
 from mcp.client.client import Client
 from mcp.server import Server, ServerRequestContext
@@ -175,6 +175,21 @@ async def test_read_resource(app: MCPServer):
         )
 
 
+async def test_read_resource_error_propagates():
+    """MCPError raised by a server handler propagates to the client with its code intact."""
+
+    async def handle_read_resource(
+        ctx: ServerRequestContext, params: types.ReadResourceRequestParams
+    ) -> ReadResourceResult:
+        raise MCPError(code=404, message="no resource with that URI was found")
+
+    server = Server("test", on_read_resource=handle_read_resource)
+    async with Client(server) as client:
+        with pytest.raises(MCPError) as exc_info:
+            await client.read_resource("unknown://example")
+        assert exc_info.value.error.code == 404
+
+
 async def test_get_prompt(app: MCPServer):
     """Test getting a prompt."""
     async with Client(app) as client:
diff --git a/tests/shared/test_ws.py b/tests/shared/test_ws.py
@@ -1,210 +1,51 @@
-import multiprocessing
-import socket
-from collections.abc import AsyncGenerator, Generator
-from urllib.parse import urlparse
+"""Smoke test for the WebSocket transport.
+
+Runs the full WS stack end-to-end over a real TCP connection, covering both
+``src/mcp/client/websocket.py`` and ``src/mcp/server/websocket.py``. MCP
+semantics (error propagation, timeouts, etc.) are transport-agnostic and are
+covered in ``tests/client/test_client.py`` and ``tests/issues/test_88_random_error.py``.
+"""
+
+from collections.abc import Generator
 
-import anyio
 import pytest
-import uvicorn
 from starlette.applications import Starlette
 from starlette.routing import WebSocketRoute
 from starlette.websockets import WebSocket
 
-from mcp import MCPError
 from mcp.client.session import ClientSession
 from mcp.client.websocket import websocket_client
-from mcp.server import Server, ServerRequestContext
+from mcp.server import Server
 from mcp.server.websocket import websocket_server
-from mcp.types import (
-    CallToolRequestParams,
-    CallToolResult,
-    EmptyResult,
-    InitializeResult,
-    ListToolsResult,
-    PaginatedRequestParams,
-    ReadResourceRequestParams,
-    ReadResourceResult,
-    TextContent,
-    TextResourceContents,
-    Tool,
-)
-from tests.test_helpers import wait_for_server
+from mcp.types import EmptyResult, InitializeResult
+from tests.test_helpers import run_uvicorn_in_thread
 
 SERVER_NAME = "test_server_for_WS"
 
 
-@pytest.fixture
-def server_port() -> int:
-    with socket.socket() as s:
-        s.bind(("127.0.0.1", 0))
-        return s.getsockname()[1]
-
-
-@pytest.fixture
-def server_url(server_port: int) -> str:
-    return f"ws://127.0.0.1:{server_port}"
-
-
-async def handle_read_resource(  # pragma: no cover
-    ctx: ServerRequestContext, params: ReadResourceRequestParams
-) -> ReadResourceResult:
-    parsed = urlparse(str(params.uri))
-    if parsed.scheme == "foobar":
-        return ReadResourceResult(
-            contents=[TextResourceContents(uri=str(params.uri), text=f"Read {parsed.netloc}", mime_type="text/plain")]
-        )
-    elif parsed.scheme == "slow":
-        await anyio.sleep(2.0)
-        return ReadResourceResult(
-            contents=[
-                TextResourceContents(
-                    uri=str(params.uri), text=f"Slow response from {parsed.netloc}", mime_type="text/plain"
-                )
-            ]
-        )
-    raise MCPError(code=404, message="OOPS! no resource with that URI was found")
-
-
-async def handle_list_tools(  # pragma: no cover
-    ctx: ServerRequestContext, params: PaginatedRequestParams | None
-) -> ListToolsResult:
-    return ListToolsResult(
-        tools=[
-            Tool(
-                name="test_tool",
-                description="A test tool",
-                input_schema={"type": "object", "properties": {}},
-            )
-        ]
-    )
-
+def make_server_app() -> Starlette:
+    srv = Server(SERVER_NAME)
 
-async def handle_call_tool(  # pragma: no cover
-    ctx: ServerRequestContext, params: CallToolRequestParams
-) -> CallToolResult:
-    return CallToolResult(content=[TextContent(type="text", text=f"Called {params.name}")])
-
-
-def _create_server() -> Server:  # pragma: no cover
-    return Server(
-        SERVER_NAME,
-        on_read_resource=handle_read_resource,
-        on_list_tools=handle_list_tools,
-        on_call_tool=handle_call_tool,
-    )
-
-
-# Test fixtures
-def make_server_app() -> Starlette:  # pragma: no cover
-    """Create test Starlette app with WebSocket transport"""
-    server = _create_server()
-
-    async def handle_ws(websocket: WebSocket):
+    async def handle_ws(websocket: WebSocket) -> None:
         async with websocket_server(websocket.scope, websocket.receive, websocket.send) as streams:
-            await server.run(streams[0], streams[1], server.create_initialization_options())
-
-    app = Starlette(routes=[WebSocketRoute("/ws", endpoint=handle_ws)])
-    return app
-
-
-def run_server(server_port: int) -> None:  # pragma: no cover
-    app = make_server_app()
-    server = uvicorn.Server(config=uvicorn.Config(app=app, host="127.0.0.1", port=server_port, log_level="error"))
-    print(f"starting server on {server_port}")
-    server.run()
-
-
-@pytest.fixture()
-def server(server_port: int) -> Generator[None, None, None]:
-    proc = multiprocessing.Process(target=run_server, kwargs={"server_port": server_port}, daemon=True)
-    print("starting process")
-    proc.start()
+            await srv.run(streams[0], streams[1], srv.create_initialization_options())
 
-    # Wait for server to be running
-    print("waiting for server to start")
-    wait_for_server(server_port)
+    return Starlette(routes=[WebSocketRoute("/ws", endpoint=handle_ws)])
 
-    yield
 
-    print("killing server")
-    # Signal the server to stop
-    proc.kill()
-    proc.join(timeout=2)
-    if proc.is_alive():  # pragma: no cover
-        print("server process failed to terminate")
-
-
-@pytest.fixture()
-async def initialized_ws_client_session(server: None, server_url: str) -> AsyncGenerator[ClientSession, None]:
-    """Create and initialize a WebSocket client session"""
-    async with websocket_client(server_url + "/ws") as streams:
-        async with ClientSession(*streams) as session:
-            # Test initialization
-            result = await session.initialize()
-            assert isinstance(result, InitializeResult)
-            assert result.server_info.name == SERVER_NAME
-
-            # Test ping
-            ping_result = await session.send_ping()
-            assert isinstance(ping_result, EmptyResult)
-
-            yield session
+@pytest.fixture
+def ws_server_url() -> Generator[str, None, None]:
+    with run_uvicorn_in_thread(make_server_app()) as base_url:
+        yield base_url.replace("http://", "ws://") + "/ws"
 
 
-# Tests
 @pytest.mark.anyio
-async def test_ws_client_basic_connection(server: None, server_url: str) -> None:
-    """Test the WebSocket connection establishment"""
-    async with websocket_client(server_url + "/ws") as streams:
+async def test_ws_client_basic_connection(ws_server_url: str) -> None:
+    async with websocket_client(ws_server_url) as streams:
         async with ClientSession(*streams) as session:
-            # Test initialization
             result = await session.initialize()
             assert isinstance(result, InitializeResult)
             assert result.server_info.name == SERVER_NAME
 
-            # Test ping
             ping_result = await session.send_ping()
             assert isinstance(ping_result, EmptyResult)
-
-
-@pytest.mark.anyio
-async def test_ws_client_happy_request_and_response(
-    initialized_ws_client_session: ClientSession,
-) -> None:
-    """Test a successful request and response via WebSocket"""
-    result = await initialized_ws_client_session.read_resource("foobar://example")
-    assert isinstance(result, ReadResourceResult)
-    assert isinstance(result.contents, list)
-    assert len(result.contents) > 0
-    assert isinstance(result.contents[0], TextResourceContents)
-    assert result.contents[0].text == "Read example"
-
-
-@pytest.mark.anyio
-async def test_ws_client_exception_handling(
-    initialized_ws_client_session: ClientSession,
-) -> None:
-    """Test exception handling in WebSocket communication"""
-    with pytest.raises(MCPError) as exc_info:
-        await initialized_ws_client_session.read_resource("unknown://example")
-    assert exc_info.value.error.code == 404
-
-
-@pytest.mark.anyio
-async def test_ws_client_timeout(
-    initialized_ws_client_session: ClientSession,
-) -> None:
-    """Test timeout handling in WebSocket communication"""
-    # Set a very short timeout to trigger a timeout exception
-    with pytest.raises(TimeoutError):
-        with anyio.fail_after(0.1):  # 100ms timeout
-            await initialized_ws_client_session.read_resource("slow://example")
-
-    # Now test that we can still use the session after a timeout
-    with anyio.fail_after(5):  # Longer timeout to allow completion
-        result = await initialized_ws_client_session.read_resource("foobar://example")
-        assert isinstance(result, ReadResourceResult)
-        assert isinstance(result.contents, list)
-        assert len(result.contents) > 0
-        assert isinstance(result.contents[0], TextResourceContents)
-        assert result.contents[0].text == "Read example"
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
@@ -1,7 +1,61 @@
 """Common test utilities for MCP server tests."""
 
 import socket
+import threading
 import time
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+import uvicorn
+
+_SERVER_SHUTDOWN_TIMEOUT_S = 5.0
+
+
+@contextmanager
+def run_uvicorn_in_thread(app: Any, **config_kwargs: Any) -> Generator[str, None, None]:
+    """Run a uvicorn server in a background thread on an ephemeral port.
+
+    The socket is bound and put into listening state *before* the thread
+    starts, so the port is known immediately with no wait. The kernel's
+    listen queue buffers any connections that arrive before uvicorn's event
+    loop reaches ``accept()``, so callers can connect as soon as this
+    function yields — no polling, no sleeps, no startup race.
+
+    This also avoids the TOCTOU race of the old pick-a-port-then-rebind
+    pattern: the socket passed here is the one uvicorn serves on, with no
+    gap where another pytest-xdist worker could claim it.
+
+    Args:
+        app: ASGI application to serve.
+        **config_kwargs: Additional keyword arguments for :class:`uvicorn.Config`
+            (e.g. ``log_level``). ``host``/``port`` are ignored since the
+            socket is pre-bound.
+
+    Yields:
+        The base URL of the running server, e.g. ``http://127.0.0.1:54321``.
+    """
+    host = "127.0.0.1"
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    sock.bind((host, 0))
+    sock.listen()
+    port = sock.getsockname()[1]
+
+    config_kwargs.setdefault("log_level", "error")
+    # Uvicorn's interface autodetection calls asyncio.iscoroutinefunction,
+    # which Python 3.14 deprecates. Under filterwarnings=error this crashes
+    # the server thread silently. Starlette is asgi3; skip the autodetect.
+    config_kwargs.setdefault("interface", "asgi3")
+    server = uvicorn.Server(config=uvicorn.Config(app=app, **config_kwargs))
+
+    thread = threading.Thread(target=server.run, kwargs={"sockets": [sock]}, daemon=True)
+    thread.start()
+    try:
+        yield f"http://{host}:{port}"
+    finally:
+        server.should_exit = True
+        thread.join(timeout=_SERVER_SHUTDOWN_TIMEOUT_S)
 
 
 def wait_for_server(port: int, timeout: float = 20.0) -> None:
