modelcontextprotocol--python-sdk/tests/shared/test_ws.py

import multiprocessing
import socket
import time
from collections.abc import AsyncGenerator, Generator
from typing import Any

import anyio
import pytest
import uvicorn
from pydantic import AnyUrl
from starlette.applications import Starlette
from starlette.routing import WebSocketRoute
from starlette.types import Message
from starlette.websockets import WebSocket

from mcp.client.session import ClientSession
from mcp.client.websocket import websocket_client
from mcp.server import Server
from mcp.server.websocket import websocket_server
from mcp.shared.exceptions import McpError
from mcp.types import (
    EmptyResult,
    ErrorData,
    InitializeResult,
    ReadResourceResult,
    TextContent,
    TextResourceContents,
    Tool,
)
from tests.test_helpers import wait_for_server

SERVER_NAME = "test_server_for_WS"

# This suite intentionally exercises the deprecated WebSocket transport.
pytestmark = pytest.mark.filterwarnings(
    "ignore:The WebSocket (client|server) transport is deprecated:DeprecationWarning"
)


@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}"


# Test server implementation
class ServerTest(Server):  # pragma: no cover
    def __init__(self):
        super().__init__(SERVER_NAME)

        @self.read_resource()
        async def handle_read_resource(uri: AnyUrl) -> str | bytes:
            if uri.scheme == "foobar":
                return f"Read {uri.host}"
            elif uri.scheme == "slow":
                # Simulate a slow resource
                await anyio.sleep(2.0)
                return f"Slow response from {uri.host}"

            raise McpError(error=ErrorData(code=404, message="OOPS! no resource with that URI was found"))

        @self.list_tools()
        async def handle_list_tools() -> list[Tool]:
            return [
                Tool(
                    name="test_tool",
                    description="A test tool",
                    inputSchema={"type": "object", "properties": {}},
                )
            ]

        @self.call_tool()
        async def handle_call_tool(name: str, args: dict[str, Any]) -> list[TextContent]:
            return [TextContent(type="text", text=f"Called {name}")]


# Test fixtures
def make_server_app() -> Starlette:  # pragma: no cover
    """Create test Starlette app with WebSocket transport"""
    server = ServerTest()

    async def handle_ws(websocket: WebSocket):
        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()

    # Give server time to start
    while not server.started:
        print("waiting for server to start")
        time.sleep(0.5)


@pytest.fixture()  # pragma: no cover
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()

    # Wait for server to be running
    print("waiting for server to start")
    wait_for_server(server_port)

    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.serverInfo.name == SERVER_NAME

            # Test ping
            ping_result = await session.send_ping()
            assert isinstance(ping_result, EmptyResult)

            yield session


# 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 with ClientSession(*streams) as session:
            # Test initialization
            result = await session.initialize()
            assert isinstance(result, InitializeResult)
            assert result.serverInfo.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(AnyUrl("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(AnyUrl("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(AnyUrl("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(AnyUrl("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"


def test_websocket_client_is_deprecated() -> None:
    """Creating the websocket_client context manager emits a DeprecationWarning."""
    with pytest.warns(DeprecationWarning, match="The WebSocket client transport is deprecated"):
        websocket_client("ws://127.0.0.1:1/ws")


def test_websocket_server_is_deprecated() -> None:
    """Creating the websocket_server context manager emits a DeprecationWarning."""

    async def receive() -> Message:
        raise NotImplementedError

    async def send(message: Message) -> None:
        raise NotImplementedError

    with pytest.warns(DeprecationWarning, match="The WebSocket server transport is deprecated"):
        websocket_server({"type": "websocket"}, receive, send)