test_helpers.py

"""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:
    """Wait for server to be ready to accept connections.

    Polls the server port until it accepts connections or timeout is reached.
    This eliminates race conditions without arbitrary sleeps.

    Args:
        port: The port number to check
        timeout: Maximum time to wait in seconds (default 5.0)

    Raises:
        TimeoutError: If server doesn't start within the timeout period
    """
    start_time = time.time()
    while time.time() - start_time < timeout:
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(0.1)
                s.connect(("127.0.0.1", port))
                # Server is ready
                return
        except (ConnectionRefusedError, OSError):
            # Server not ready yet, retry quickly
            time.sleep(0.01)
    raise TimeoutError(f"Server on port {port} did not start within {timeout} seconds")  # pragma: no cover