Fixed flaky server-side SSE tests

Author: agronholm

tests/server/test_sse_security.py

@@ -4,6 +4,7 @@
 import multiprocessing
 import socket
 import time
+from contextlib import contextmanager
 
 import httpx
 import pytest
@@ -66,40 +67,48 @@ async def handle_sse(request: Request):
     uvicorn.run(starlette_app, host="127.0.0.1", port=port, log_level="error")
 
 
+@contextmanager
 def start_server_process(port: int, security_settings: TransportSecuritySettings | None = None):
     """Start server in a separate process."""
     context = multiprocessing.get_context("spawn")
     process = context.Process(target=run_server_with_settings, args=(port, security_settings))
     process.start()
-    # Give server time to start
-    time.sleep(1)
-    return process
+
+    # Wait until the designated port can be connected
+    max_attempts = 20
+    for attempt in range(max_attempts):
+        try:
+            with socket.create_connection(("127.0.0.1", port)):
+                break
+        except ConnectionRefusedError:
+            time.sleep(0.1)
+    else:
+        raise RuntimeError(f"Server failed to start after {max_attempts} attempts")
+
+    try:
+        yield
+    finally:
+        process.terminate()
+        process.join()
 
 
 @pytest.mark.anyio
 async def test_sse_security_default_settings(server_port: int):
     """Test SSE with default security settings (protection disabled)."""
-    process = start_server_process(server_port)
-
-    try:
+    with start_server_process(server_port):
         headers = {"Host": "evil.com", "Origin": "http://evil.com"}
 
         async with httpx.AsyncClient(timeout=5.0) as client:
             async with client.stream("GET", f"http://127.0.0.1:{server_port}/sse", headers=headers) as response:
                 assert response.status_code == 200
-    finally:
-        process.terminate()
-        process.join()
 
 
 @pytest.mark.anyio
 async def test_sse_security_invalid_host_header(server_port: int):
     """Test SSE with invalid Host header."""
     # Enable security by providing settings with an empty allowed_hosts list
     security_settings = TransportSecuritySettings(enable_dns_rebinding_protection=True, allowed_hosts=["example.com"])
-    process = start_server_process(server_port, security_settings)
-
-    try:
+    with start_server_process(server_port, security_settings):
         # Test with invalid host header
         headers = {"Host": "evil.com"}
 
@@ -108,10 +117,6 @@ async def test_sse_security_invalid_host_header(server_port: int):
             assert response.status_code == 421
             assert response.text == "Invalid Host header"
 
-    finally:
-        process.terminate()
-        process.join()
-
 
 @pytest.mark.anyio
 async def test_sse_security_invalid_origin_header(server_port: int):
@@ -120,9 +125,7 @@ async def test_sse_security_invalid_origin_header(server_port: int):
     security_settings = TransportSecuritySettings(
         enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1:*"], allowed_origins=["http://localhost:*"]
     )
-    process = start_server_process(server_port, security_settings)
-
-    try:
+    with start_server_process(server_port, security_settings):
         # Test with invalid origin header
         headers = {"Origin": "http://evil.com"}
 
@@ -131,10 +134,6 @@ async def test_sse_security_invalid_origin_header(server_port: int):
             assert response.status_code == 400
             assert response.text == "Invalid Origin header"
 
-    finally:
-        process.terminate()
-        process.join()
-
 
 @pytest.mark.anyio
 async def test_sse_security_post_invalid_content_type(server_port: int):
@@ -143,9 +142,7 @@ async def test_sse_security_post_invalid_content_type(server_port: int):
     security_settings = TransportSecuritySettings(
         enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1:*"], allowed_origins=["http://127.0.0.1:*"]
     )
-    process = start_server_process(server_port, security_settings)
-
-    try:
+    with start_server_process(server_port, security_settings):
         async with httpx.AsyncClient(timeout=5.0) as client:
             # Test POST with invalid content type
             fake_session_id = "12345678123456781234567812345678"
@@ -164,18 +161,12 @@ async def test_sse_security_post_invalid_content_type(server_port: int):
             assert response.status_code == 400
             assert response.text == "Invalid Content-Type header"
 
-    finally:
-        process.terminate()
-        process.join()
-
 
 @pytest.mark.anyio
 async def test_sse_security_disabled(server_port: int):
     """Test SSE with security disabled."""
     settings = TransportSecuritySettings(enable_dns_rebinding_protection=False)
-    process = start_server_process(server_port, settings)
-
-    try:
+    with start_server_process(server_port, settings):
         # Test with invalid host header - should still work
         headers = {"Host": "evil.com"}
 
@@ -185,10 +176,6 @@ async def test_sse_security_disabled(server_port: int):
                 # Should connect successfully even with invalid host
                 assert response.status_code == 200
 
-    finally:
-        process.terminate()
-        process.join()
-
 
 @pytest.mark.anyio
 async def test_sse_security_custom_allowed_hosts(server_port: int):
@@ -198,9 +185,7 @@ async def test_sse_security_custom_allowed_hosts(server_port: int):
         allowed_hosts=["localhost", "127.0.0.1", "custom.host"],
         allowed_origins=["http://localhost", "http://127.0.0.1", "http://custom.host"],
     )
-    process = start_server_process(server_port, settings)
-
-    try:
+    with start_server_process(server_port, settings):
         # Test with custom allowed host
         headers = {"Host": "custom.host"}
 
@@ -218,10 +203,6 @@ async def test_sse_security_custom_allowed_hosts(server_port: int):
             assert response.status_code == 421
             assert response.text == "Invalid Host header"
 
-    finally:
-        process.terminate()
-        process.join()
-
 
 @pytest.mark.anyio
 async def test_sse_security_wildcard_ports(server_port: int):
@@ -231,9 +212,7 @@ async def test_sse_security_wildcard_ports(server_port: int):
         allowed_hosts=["localhost:*", "127.0.0.1:*"],
         allowed_origins=["http://localhost:*", "http://127.0.0.1:*"],
     )
-    process = start_server_process(server_port, settings)
-
-    try:
+    with start_server_process(server_port, settings):
         # Test with various port numbers
         for test_port in [8080, 3000, 9999]:
             headers = {"Host": f"localhost:{test_port}"}
@@ -252,10 +231,6 @@ async def test_sse_security_wildcard_ports(server_port: int):
                     # Should connect successfully with any port
                     assert response.status_code == 200
 
-    finally:
-        process.terminate()
-        process.join()
-
 
 @pytest.mark.anyio
 async def test_sse_security_post_valid_content_type(server_port: int):
@@ -264,9 +239,7 @@ async def test_sse_security_post_valid_content_type(server_port: int):
     security_settings = TransportSecuritySettings(
         enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1:*"], allowed_origins=["http://127.0.0.1:*"]
     )
-    process = start_server_process(server_port, security_settings)
-
-    try:
+    with start_server_process(server_port, security_settings):
         async with httpx.AsyncClient() as client:
             # Test with various valid content types
             valid_content_types = [
@@ -288,7 +261,3 @@ async def test_sse_security_post_valid_content_type(server_port: int):
                 # We're testing that it passes the content-type check
                 assert response.status_code == 404
                 assert response.text == "Could not find session"
-
-    finally:
-        process.terminate()
-        process.join()