fix: Replace remaining manual server polling with wait_for_server helper (modelcontextprotocol#1529)

Author: maxisbey

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

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")

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
 

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