diff --git a/tests/client/test_http_unicode.py b/tests/client/test_http_unicode.py index edf8675e5..a0737e01d 100644 --- a/tests/client/test_http_unicode.py +++ b/tests/client/test_http_unicode.py @@ -7,13 +7,13 @@ import multiprocessing import socket -import time from collections.abc import Generator import pytest from mcp.client.session import ClientSession from mcp.client.streamable_http import streamablehttp_client +from tests.test_helpers import wait_for_server # Test constants with various Unicode characters UNICODE_TEST_STRINGS = { @@ -158,19 +158,8 @@ def running_unicode_server(unicode_server_port: int) -> Generator[str, None, Non proc = multiprocessing.Process(target=run_unicode_server, kwargs={"port": unicode_server_port}, daemon=True) 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", unicode_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 to be ready + wait_for_server(unicode_server_port) try: yield f"http://127.0.0.1:{unicode_server_port}" diff --git a/tests/client/test_notification_response.py b/tests/client/test_notification_response.py index 88e64711b..3840d4262 100644 --- a/tests/client/test_notification_response.py +++ b/tests/client/test_notification_response.py @@ -8,7 +8,6 @@ import json import multiprocessing import socket -import time from collections.abc import Generator import pytest @@ -22,6 +21,7 @@ from mcp.client.streamable_http import streamablehttp_client from mcp.shared.session import RequestResponder from mcp.types import ClientNotification, RootsListChangedNotification +from tests.test_helpers import wait_for_server def create_non_sdk_server_app() -> Starlette: @@ -95,14 +95,9 @@ def non_sdk_server(non_sdk_server_port: int) -> Generator[None, None, None]: proc.start() # Wait for server to be ready - start_time = time.time() - while time.time() - start_time < 10: - try: - with socket.create_connection(("127.0.0.1", non_sdk_server_port), timeout=0.1): - break - except (TimeoutError, ConnectionRefusedError): - time.sleep(0.1) - else: + try: + wait_for_server(non_sdk_server_port, timeout=10.0) + except TimeoutError: proc.kill() proc.join(timeout=2) pytest.fail("Server failed to start within 10 seconds") diff --git a/tests/server/fastmcp/test_integration.py b/tests/server/fastmcp/test_integration.py index dc88cc025..618d7bc61 100644 --- a/tests/server/fastmcp/test_integration.py +++ b/tests/server/fastmcp/test_integration.py @@ -13,7 +13,6 @@ import json import multiprocessing import socket -import time from collections.abc import Generator import pytest @@ -60,6 +59,7 @@ TextResourceContents, ToolListChangedNotification, ) +from tests.test_helpers import wait_for_server class NotificationCollector: @@ -160,19 +160,8 @@ def server_transport(request: pytest.FixtureRequest, server_port: int) -> Genera ) 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", 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 to be ready + wait_for_server(server_port) yield transport diff --git a/tests/shared/test_sse.py b/tests/shared/test_sse.py index 17847497f..fdb6ccfd8 100644 --- a/tests/shared/test_sse.py +++ b/tests/shared/test_sse.py @@ -252,19 +252,8 @@ def mounted_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 @@ -367,19 +356,8 @@ def context_server(server_port: int) -> Generator[None, None, None]: proc.start() # Wait for server to be running - max_attempts = 20 - attempt = 0 print("waiting for context 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"Context server failed to start after {max_attempts} attempts") + wait_for_server(server_port) yield