diff --git a/pyproject.toml b/pyproject.toml index 21013e79e7..e0498ec52f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -98,7 +98,7 @@ venv = ".venv" # those private functions instead of testing the private functions directly. It makes it easier to maintain the code source # and refactor code that is not public. executionEnvironments = [ - { root = "tests", reportUnusedFunction = false, reportPrivateUsage = false }, + { root = "tests", extraPaths = ["."], reportUnusedFunction = false, reportPrivateUsage = false }, { root = "examples/servers", reportUnusedFunction = false }, ] diff --git a/tests/server/test_sse_security.py b/tests/server/test_sse_security.py index 5a2210b8e1..7a8e52bdab 100644 --- a/tests/server/test_sse_security.py +++ b/tests/server/test_sse_security.py @@ -3,7 +3,6 @@ import logging import multiprocessing import socket -import time import httpx import pytest @@ -17,6 +16,7 @@ from mcp.server.sse import SseServerTransport from mcp.server.transport_security import TransportSecuritySettings from mcp.types import Tool +from tests.test_helpers import wait_for_server logger = logging.getLogger(__name__) SERVER_NAME = "test_sse_security_server" @@ -66,26 +66,6 @@ async def handle_sse(request: Request): uvicorn.run(starlette_app, host="127.0.0.1", port=port, log_level="error") -def wait_for_server(port: int, timeout: float = 5.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. - """ - 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") - - def start_server_process(port: int, security_settings: TransportSecuritySettings | None = None): """Start server in a separate process.""" process = multiprocessing.Process(target=run_server_with_settings, args=(port, security_settings)) diff --git a/tests/server/test_streamable_http_security.py b/tests/server/test_streamable_http_security.py index b9cd83dc1b..de302fb7c6 100644 --- a/tests/server/test_streamable_http_security.py +++ b/tests/server/test_streamable_http_security.py @@ -3,7 +3,6 @@ import logging import multiprocessing import socket -import time from collections.abc import AsyncGenerator from contextlib import asynccontextmanager @@ -18,6 +17,7 @@ from mcp.server.streamable_http_manager import StreamableHTTPSessionManager from mcp.server.transport_security import TransportSecuritySettings from mcp.types import Tool +from tests.test_helpers import wait_for_server logger = logging.getLogger(__name__) SERVER_NAME = "test_streamable_http_security_server" @@ -77,8 +77,8 @@ def start_server_process(port: int, security_settings: TransportSecuritySettings """Start server in a separate process.""" process = multiprocessing.Process(target=run_server_with_settings, args=(port, security_settings)) process.start() - # Give server time to start - time.sleep(1) + # Wait for server to be ready to accept connections + wait_for_server(port) return process diff --git a/tests/shared/test_sse.py b/tests/shared/test_sse.py index 7b0d89cb42..17847497f9 100644 --- a/tests/shared/test_sse.py +++ b/tests/shared/test_sse.py @@ -32,6 +32,7 @@ TextResourceContents, Tool, ) +from tests.test_helpers import wait_for_server SERVER_NAME = "test_server_for_SSE" @@ -123,19 +124,8 @@ def server(server_port: int) -> Generator[None, None, None]: proc.start() # Wait for server to be running - max_attempts = 20 - attempt = 0 print("waiting for server to start") - while attempt < max_attempts: - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.connect(("127.0.0.1", server_port)) - break - except ConnectionRefusedError: - time.sleep(0.1) - attempt += 1 - else: - raise RuntimeError(f"Server failed to start after {max_attempts} attempts") + wait_for_server(server_port) yield diff --git a/tests/shared/test_streamable_http.py b/tests/shared/test_streamable_http.py index 34e9291680..0cc85f4417 100644 --- a/tests/shared/test_streamable_http.py +++ b/tests/shared/test_streamable_http.py @@ -7,7 +7,6 @@ import json import multiprocessing import socket -import time from collections.abc import Generator from typing import Any @@ -43,6 +42,7 @@ from mcp.shared.message import ClientMessageMetadata from mcp.shared.session import RequestResponder from mcp.types import InitializeResult, TextContent, TextResourceContents, Tool +from tests.test_helpers import wait_for_server # Test constants SERVER_NAME = "test_streamable_http_server" @@ -344,18 +344,7 @@ def basic_server(basic_server_port: int) -> Generator[None, None, None]: proc.start() # Wait for server to be running - max_attempts = 20 - attempt = 0 - while attempt < max_attempts: - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.connect(("127.0.0.1", basic_server_port)) - break - except ConnectionRefusedError: - time.sleep(0.1) - attempt += 1 - else: - raise RuntimeError(f"Server failed to start after {max_attempts} attempts") + wait_for_server(basic_server_port) yield @@ -391,18 +380,7 @@ def event_server( proc.start() # Wait for server to be running - max_attempts = 20 - attempt = 0 - while attempt < max_attempts: - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.connect(("127.0.0.1", event_server_port)) - break - except ConnectionRefusedError: - time.sleep(0.1) - attempt += 1 - else: - raise RuntimeError(f"Server failed to start after {max_attempts} attempts") + wait_for_server(event_server_port) yield event_store, f"http://127.0.0.1:{event_server_port}" @@ -422,18 +400,7 @@ def json_response_server(json_server_port: int) -> Generator[None, None, None]: proc.start() # Wait for server to be running - max_attempts = 20 - attempt = 0 - while attempt < max_attempts: - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.connect(("127.0.0.1", json_server_port)) - break - except ConnectionRefusedError: - time.sleep(0.1) - attempt += 1 - else: - raise RuntimeError(f"Server failed to start after {max_attempts} attempts") + wait_for_server(json_server_port) yield @@ -1407,18 +1374,7 @@ def context_aware_server(basic_server_port: int) -> Generator[None, None, None]: proc.start() # Wait for server to be running - max_attempts = 20 - attempt = 0 - while attempt < max_attempts: - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.connect(("127.0.0.1", basic_server_port)) - break - except ConnectionRefusedError: - time.sleep(0.1) - attempt += 1 - else: - raise RuntimeError(f"Context-aware server failed to start after {max_attempts} attempts") + wait_for_server(basic_server_port) yield diff --git a/tests/shared/test_ws.py b/tests/shared/test_ws.py index 2d67eccdd0..71b0d4cc00 100644 --- a/tests/shared/test_ws.py +++ b/tests/shared/test_ws.py @@ -26,6 +26,7 @@ TextResourceContents, Tool, ) +from tests.test_helpers import wait_for_server SERVER_NAME = "test_server_for_WS" @@ -110,19 +111,8 @@ def server(server_port: int) -> Generator[None, None, None]: proc.start() # Wait for server to be running - max_attempts = 20 - attempt = 0 print("waiting for server to start") - while attempt < max_attempts: - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.connect(("127.0.0.1", server_port)) - break - except ConnectionRefusedError: - time.sleep(0.1) - attempt += 1 - else: - raise RuntimeError(f"Server failed to start after {max_attempts} attempts") + wait_for_server(server_port) yield diff --git a/tests/test_helpers.py b/tests/test_helpers.py new file mode 100644 index 0000000000..a4b4146e91 --- /dev/null +++ b/tests/test_helpers.py @@ -0,0 +1,31 @@ +"""Common test utilities for MCP server tests.""" + +import socket +import time + + +def wait_for_server(port: int, timeout: float = 5.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")