|
2 | 2 |
|
3 | 3 | import socket |
4 | 4 | import time |
5 | | -import urllib.error |
6 | | -import urllib.request |
7 | 5 |
|
8 | 6 |
|
9 | 7 | def wait_for_server(port: int, timeout: float = 20.0) -> None: |
10 | 8 | """Wait for server to be ready to accept connections. |
11 | 9 |
|
12 | | - First polls until the TCP port accepts connections, then verifies the |
13 | | - HTTP server is actually ready to handle requests. This two-stage check |
14 | | - prevents race conditions where the port is open but the ASGI app hasn't |
15 | | - finished initializing. |
| 10 | + Polls the server port until it accepts connections or timeout is reached. |
| 11 | + This eliminates race conditions without arbitrary sleeps. |
16 | 12 |
|
17 | 13 | Args: |
18 | 14 | port: The port number to check |
19 | | - timeout: Maximum time to wait in seconds (default 20.0) |
| 15 | + timeout: Maximum time to wait in seconds (default 5.0) |
20 | 16 |
|
21 | 17 | Raises: |
22 | 18 | TimeoutError: If server doesn't start within the timeout period |
23 | 19 | """ |
24 | 20 | start_time = time.time() |
25 | | - |
26 | | - # Stage 1: Wait for TCP port to accept connections |
27 | 21 | while time.time() - start_time < timeout: |
28 | 22 | try: |
29 | 23 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: |
30 | 24 | s.settimeout(0.1) |
31 | 25 | s.connect(("127.0.0.1", port)) |
32 | | - break # Port is open, move to stage 2 |
| 26 | + # Server is ready |
| 27 | + return |
33 | 28 | except (ConnectionRefusedError, OSError): |
| 29 | + # Server not ready yet, retry quickly |
34 | 30 | time.sleep(0.01) |
35 | | - else: |
36 | | - raise TimeoutError(f"Server on port {port} did not start within {timeout} seconds") # pragma: no cover |
37 | | - |
38 | | - # Stage 2: Verify HTTP server is ready by making a request. |
39 | | - # A non-existent path returns 404/405 if the app is ready, or |
40 | | - # raises an error if the ASGI app hasn't finished initializing. |
41 | | - while time.time() - start_time < timeout: |
42 | | - try: |
43 | | - req = urllib.request.Request( |
44 | | - f"http://127.0.0.1:{port}/healthz", |
45 | | - method="GET", |
46 | | - ) |
47 | | - with urllib.request.urlopen(req, timeout=1): |
48 | | - return # Any successful response means server is ready |
49 | | - except urllib.error.HTTPError: |
50 | | - # 404/405/etc means the server IS handling requests |
51 | | - return |
52 | | - except (urllib.error.URLError, ConnectionError, OSError): |
53 | | - # Server not ready for HTTP yet |
54 | | - time.sleep(0.05) |
55 | | - raise TimeoutError(f"Server on port {port} did not become HTTP-ready within {timeout} seconds") # pragma: no cover |
| 31 | + raise TimeoutError(f"Server on port {port} did not start within {timeout} seconds") # pragma: no cover |
0 commit comments