Upgrade websockets and FastAPI version in tests and dev (#662)

Author: lexiforest

pyproject.toml

@@ -42,7 +42,7 @@ dev = [
     "trio>=0.25.0,<1.0",
     "trustme>=1.1.0,<2.0",
     "uvicorn>=0.29.0,<1.0",
-    "websockets>=12.0,<13.0",
+    "websockets>=14.0",
     "typing_extensions",
 ]
 build = [
@@ -52,7 +52,7 @@ build = [
 test = [
     "charset_normalizer>=3.3.2,<4.0",
     "cryptography>=42.0.5,<43.0",
-    "fastapi==0.110.0,<1.0",
+    "fastapi>=0.110.0,<1.0",
     "httpx==0.23.1",  # don't change, tests will raise "httpx.InvalidURL: Invalid URL component 'path'"
     "proxy.py>=2.4.3,<3.0",
     "pytest>=8.1.1,<9.0",
@@ -62,7 +62,7 @@ test = [
     "trio>=0.25.0,<1.0",
     "trustme>=1.1.0,<2.0",
     "uvicorn>=0.29.0,<1.0",
-    "websockets>=12.0",
+    "websockets>=14.0",
     "typing_extensions",
 ]
 

tests/unittest/conftest.py

@@ -1,7 +1,9 @@
 import asyncio
 import contextlib
+from dataclasses import dataclass
 import json
 import os
+import queue
 import threading
 import time
 import typing
@@ -15,6 +17,7 @@
 import trustme
 import uvicorn
 import websockets
+from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives.serialization import (
     BestAvailableEncryption,
@@ -613,30 +616,11 @@ async def watch_restarts(self):  # pragma: nocover
             await self.startup()
 
 
-async def echo(websocket):
-    while True:
-        try:
-            # Echo every text or binary message
-            async for message in websocket:
-                await websocket.send(message)
-
-        except websockets.exceptions.ConnectionClosed as e:
-            # Client sent us a close frame: echo it back exactly
-            await websocket.close(code=e.code, reason=e.reason)
-
-
-class TestWebsocketServer:
-    def __init__(self, port):
-        self.url = f"ws://127.0.0.1:{port}"
-        self.port = port
-
-    def run(self):
-        async def serve(port):
-            # GitHub actions only likes 127, not localhost, wtf...
-            async with websockets.serve(echo, "127.0.0.1", port):  # pyright: ignore
-                await asyncio.Future()  # run forever
-
-        asyncio.run(serve(self.port))
+@pytest.fixture(scope="session")
+def server():
+    config = Config(app=app, lifespan="off", loop="asyncio")
+    server = TestServer(config=config)
+    yield from serve_in_thread(server)
 
 
 def serve_in_thread(server: Server):
@@ -651,26 +635,6 @@ def serve_in_thread(server: Server):
         thread.join()
 
 
-@pytest.fixture(scope="session")
-def ws_server():
-    server = TestWebsocketServer(port=8964)
-    thread = threading.Thread(target=server.run, daemon=True)
-    thread.start()
-    try:
-        time.sleep(2)  # FIXME find a reliable way to check the server is up
-        yield server
-    finally:
-        pass
-        # thread.join()
-
-
-@pytest.fixture(scope="session")
-def server():
-    config = Config(app=app, lifespan="off", loop="asyncio")
-    server = TestServer(config=config)
-    yield from serve_in_thread(server)
-
-
 @pytest.fixture(scope="session")
 def https_server(cert_pem_file, cert_private_key_file):
     config = Config(
@@ -685,6 +649,72 @@ def https_server(cert_pem_file, cert_private_key_file):
     yield from serve_in_thread(server)
 
 
+async def echo(ws):
+    try:
+        async for msg in ws:
+            await ws.send(msg)
+    except (ConnectionClosedOK, ConnectionClosedError):
+        # Normal / abnormal close — nothing extra to do.
+        pass
+
+
+def start_ws_server(port: int = 8964):
+    """
+    Start a websockets server on 127.0.0.1:port in a background thread.
+    Returns (url, stop) where stop() shuts it down.
+    """
+    ready = threading.Event()
+    stop_callable_q: queue.Queue[typing.Callable] = queue.Queue()
+
+    def _thread_target():
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        stop_async = asyncio.Event()
+
+        def _stop():
+            # can be called from main thread
+            loop.call_soon_threadsafe(stop_async.set)
+
+        async def _run():
+            async with websockets.serve(echo, "127.0.0.1", port) as _:
+                stop_callable_q.put(_stop)
+                ready.set()
+                await stop_async.wait()
+
+        try:
+            loop.run_until_complete(_run())
+        finally:
+            loop.run_until_complete(loop.shutdown_asyncgens())
+            loop.close()
+
+    t = threading.Thread(target=_thread_target, daemon=True)
+    t.start()
+
+    # Wait until server is really listening and we have a stop() handle
+    stop = stop_callable_q.get()  # blocks until put()
+    ready.wait()  # the socket is bound now
+
+    url = f"ws://127.0.0.1:{port}"
+    return url, stop, t
+
+
+@dataclass
+class WSServer:
+    url: str
+    stop: typing.Callable
+
+
+@pytest.fixture(scope="session")
+def ws_server():
+    url, stop, thread = start_ws_server(port=8964)
+    try:
+        yield WSServer(url=url, stop=stop)
+    finally:
+        stop()  # trigger graceful shutdown
+        thread.join(5)  # optional: wait up to 5s for thread to exit
+
+
 @pytest.fixture(scope="session")
 def proxy_server(request):
     ps = proxy.Proxy(port=8002, plugins=["proxy.plugin.ManInTheMiddlePlugin"])