From 5dab5e9eef8b153ca6492e746d627d0ce30208d0 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 15:21:20 -0600 Subject: [PATCH 1/9] test: add a scriptable fake Massive WebSocket server FakeServer listens on a free port and records every client message per connection. It speaks the Massive protocol (connected, auth_success or auth_failed, per-topic success acks) and can be scripted: replace the auth reply, reject topics or handshakes, send status and data messages, and close or drop connections with a chosen code. Tests cover connect, auth, subscribe, reconnect, and independent clients in parallel threads. Co-Authored-By: Claude Opus 5.5 --- test_websocket/fake_server.py | 342 +++++++++++++++++++++++++++++ test_websocket/test_fake_server.py | 164 ++++++++++++++ 2 files changed, 506 insertions(+) create mode 100644 test_websocket/fake_server.py create mode 100644 test_websocket/test_fake_server.py diff --git a/test_websocket/fake_server.py b/test_websocket/fake_server.py new file mode 100644 index 00000000..bf77a4b8 --- /dev/null +++ b/test_websocket/fake_server.py @@ -0,0 +1,342 @@ +""" +A scriptable fake Massive WebSocket server for tests. + +Unlike ``mock_server``, which listens on a fixed port and keeps module-level +state, each ``FakeServer`` listens on a free port and records everything per +instance, so tests can run several servers and assert on exactly what a client +sent. Run it inside the test's event loop:: + + async with FakeServer() as server: + c = WebSocketClient(api_key="key", feed=server.feed, secure=False) + c.subscribe("T.AAPL") + task = asyncio.create_task(c.connect(processor)) + await server.wait_for_subscriptions({"T.AAPL"}) + await server.send_data([{"ev": "T", "sym": "AAPL", "p": 1.0}]) + await server.close_connections(code=1011) # simulate an outage + +By default the server speaks the Massive protocol: it sends ``connected``, +answers ``auth`` with ``auth_success`` (or ``auth_failed`` for a key in +``reject_keys``), and acknowledges each subscribed or unsubscribed topic with a +``success`` status. Scripted replies and hooks override that behaviour. +""" + +import asyncio +import json +from http import HTTPStatus +from typing import ( + Any, + Awaitable, + Callable, + Dict, + Iterable, + List, + Optional, + Set, +) + +from websockets.asyncio.server import Server, ServerConnection, serve +from websockets.exceptions import ConnectionClosed +from websockets.http11 import Request, Response + + +def status(status: str, message: str) -> Dict[str, Any]: + """Build a Massive status message, e.g. ``status("max_connections", "...")``.""" + return {"ev": "status", "status": status, "message": message} + + +class AuthReply: + """ + A scripted reply to one ``auth`` message. + + :param frames: The messages to send, as a JSON array. + :param close_code: If set, close the connection with this code afterwards. + :param close_reason: The reason sent with the close frame. + """ + + def __init__( + self, + frames: List[Dict[str, Any]], + close_code: Optional[int] = None, + close_reason: str = "", + ): + self.frames = frames + self.close_code = close_code + self.close_reason = close_reason + + +class FakeConnection: + """One client connection accepted by a ``FakeServer``.""" + + def __init__(self, server: "FakeServer", websocket: ServerConnection, index: int): + self.server = server + self.websocket = websocket + self.index = index + self.path = websocket.request.path if websocket.request else "" + self.api_key: Optional[str] = None + self.authenticated = False + # every message the client sent, decoded from JSON + self.inbox: List[Any] = [] + # topics the client is subscribed to on this connection + self.subscriptions: Set[str] = set() + self.closed = asyncio.Event() + + @property + def open(self) -> bool: + return not self.closed.is_set() + + async def send(self, frames: List[Dict[str, Any]]): + """Send a JSON array of messages to this client.""" + await self.websocket.send(json.dumps(frames)) + + async def send_raw(self, message: str): + """Send a message as-is.""" + await self.websocket.send(message) + + async def close(self, code: int = 1000, reason: str = ""): + """Close with a close frame; 1000 is a clean close, e.g. 1011 is not.""" + await self.websocket.close(code, reason) + await self.closed.wait() + + async def drop(self): + """Drop the TCP connection without a close frame.""" + transport = self.websocket.transport + if transport is not None: + transport.abort() + await self.closed.wait() + + +class FakeServer: + """ + A fake Massive WebSocket server listening on ``host`` and a free port. + + :param host: The interface to listen on. + :param valid_keys: If set, only these API keys authenticate. + :param reject_keys: API keys answered with ``auth_failed``. + :param ack_subscriptions: Whether to acknowledge subscribe/unsubscribe. + :param reject_topics: Topics answered with an ``error`` status instead of + being subscribed. + :param on_message: Called with every decoded client message after it has + been recorded and handled; use it to script replies. + """ + + def __init__( + self, + host: str = "127.0.0.1", + valid_keys: Optional[Iterable[str]] = None, + reject_keys: Iterable[str] = (), + ack_subscriptions: bool = True, + reject_topics: Iterable[str] = (), + on_message: Optional[Callable[[FakeConnection, Any], Awaitable[None]]] = None, + ): + self.host = host + self.port = 0 + self.valid_keys = set(valid_keys) if valid_keys is not None else None + self.reject_keys = set(reject_keys) + self.ack_subscriptions = ack_subscriptions + self.reject_topics = set(reject_topics) + self.on_message = on_message + # while True, new connections get an HTTP 503 instead of a handshake + self.reject_handshakes = False + self.connections: List[FakeConnection] = [] + self._auth_replies: List[AuthReply] = [] + self._server: Optional[Server] = None + + # lifecycle + + async def start(self) -> "FakeServer": + self._server = await serve( + self._handler, + self.host, + self.port, + process_request=self._process_request, + ) + self.port = next(iter(self._server.sockets)).getsockname()[1] + return self + + async def stop(self): + if self._server is not None: + self._server.close() + await self._server.wait_closed() + self._server = None + + async def __aenter__(self) -> "FakeServer": + return await self.start() + + async def __aexit__(self, *exc): + await self.stop() + + @property + def feed(self) -> str: + """The ``feed`` to pass to ``WebSocketClient`` (with ``secure=False``).""" + return f"{self.host}:{self.port}" + + # scripting + + def script_auth( + self, + *frames: Dict[str, Any], + close_code: Optional[int] = None, + close_reason: str = "", + ): + """ + Queue the reply to the next ``auth`` message, replacing the default. + + Replies are used in order, one per connection, e.g. + ``script_auth(status("max_connections", "..."), close_code=1008)``. + """ + self._auth_replies.append(AuthReply(list(frames), close_code, close_reason)) + + @property + def connection(self) -> FakeConnection: + """The most recent connection.""" + if not self.connections: + raise AssertionError("no client has connected") + return self.connections[-1] + + @property + def open_connections(self) -> List[FakeConnection]: + return [c for c in self.connections if c.open] + + @property + def inbox(self) -> List[Any]: + """Every message received, from all connections, in order per connection.""" + return [m for c in self.connections for m in c.inbox] + + @property + def subscriptions(self) -> Set[str]: + """The topics subscribed on the most recent connection.""" + return set(self.connections[-1].subscriptions) if self.connections else set() + + async def send(self, frames: List[Dict[str, Any]]): + """Send messages to every open, authenticated connection.""" + for c in self.open_connections: + if c.authenticated: + await c.send(frames) + + async def send_status(self, status_: str, message: str): + await self.send([status(status_, message)]) + + async def send_data(self, events: List[Dict[str, Any]]): + """Send data events, e.g. ``[{"ev": "T", "sym": "AAPL", ...}]``.""" + await self.send(events) + + async def close_connections(self, code: int = 1000, reason: str = ""): + """Close every open connection with ``code``.""" + await asyncio.gather(*(c.close(code, reason) for c in self.open_connections)) + + async def drop_connections(self): + """Drop every open connection without a close frame.""" + await asyncio.gather(*(c.drop() for c in self.open_connections)) + + # waiting + + async def wait_for(self, predicate: Callable[[], bool], timeout: float = 5): + """Wait until ``predicate()`` is true, polling every 10 ms.""" + + async def wait(): + while not predicate(): + await asyncio.sleep(0.01) + + await asyncio.wait_for(wait(), timeout) + + async def wait_for_connections(self, n: int, timeout: float = 5): + """Wait until ``n`` clients have authenticated, in total.""" + await self.wait_for( + lambda: sum(c.authenticated for c in self.connections) >= n, timeout + ) + + async def wait_for_subscriptions(self, topics: Set[str], timeout: float = 5): + """Wait until the most recent connection is subscribed to exactly ``topics``.""" + await self.wait_for( + lambda: bool(self.connections) and self.subscriptions == set(topics), + timeout, + ) + + async def wait_for_message( + self, predicate: Callable[[Any], bool], timeout: float = 5 + ) -> Any: + """Wait for a client message matching ``predicate`` and return it.""" + found: List[Any] = [] + + def check() -> bool: + found[:] = [m for m in self.inbox if predicate(m)] + return bool(found) + + await self.wait_for(check, timeout) + return found[0] + + # protocol + + def _process_request( + self, websocket: ServerConnection, request: Request + ) -> Optional[Response]: + if self.reject_handshakes: + return websocket.respond(HTTPStatus.SERVICE_UNAVAILABLE, "unavailable\n") + return None + + async def _handler(self, websocket: ServerConnection): + conn = FakeConnection(self, websocket, len(self.connections)) + self.connections.append(conn) + try: + await conn.send([status("connected", "Connected Successfully")]) + async for raw in websocket: + try: + msg = json.loads(raw) + except ValueError: + msg = raw + conn.inbox.append(msg) + if isinstance(msg, dict): + await self._handle(conn, msg) + if self.on_message is not None: + await self.on_message(conn, msg) + except ConnectionClosed: + pass + finally: + conn.closed.set() + + async def _handle(self, conn: FakeConnection, msg: Dict[str, Any]): + action = msg.get("action") + params = msg.get("params") + if action == "auth": + await self._auth(conn, params) + elif action in ("subscribe", "unsubscribe") and isinstance(params, str): + topics = [p.strip() for p in params.split(",") if p.strip()] + await self._subscription(conn, action, topics) + + async def _auth(self, conn: FakeConnection, key: Optional[str]): + conn.api_key = key + if self._auth_replies: + reply = self._auth_replies.pop(0) + conn.authenticated = any( + f.get("status") == "auth_success" for f in reply.frames + ) + if reply.frames: + await conn.send(reply.frames) + if reply.close_code is not None: + await conn.websocket.close(reply.close_code, reply.close_reason) + return + if key in self.reject_keys or ( + self.valid_keys is not None and key not in self.valid_keys + ): + await conn.send([status("auth_failed", "authentication failed")]) + await conn.websocket.close(1008, "authentication failed") + return + conn.authenticated = True + await conn.send([status("auth_success", "authenticated")]) + + async def _subscription(self, conn: FakeConnection, action: str, topics: List[str]): + replies = [] + for t in topics: + if action == "subscribe" and t in self.reject_topics: + replies.append(status("error", f"subscription rejected: {t}")) + continue + if action == "subscribe": + conn.subscriptions.add(t) + ack = f"subscribed to: {t}" + else: + conn.subscriptions.discard(t) + ack = f"unsubscribed to: {t}" + if self.ack_subscriptions: + replies.append(status("success", ack)) + if replies: + await conn.send(replies) diff --git a/test_websocket/test_fake_server.py b/test_websocket/test_fake_server.py new file mode 100644 index 00000000..ede70ac3 --- /dev/null +++ b/test_websocket/test_fake_server.py @@ -0,0 +1,164 @@ +import asyncio +import threading +import unittest +from concurrent.futures import ThreadPoolExecutor +from typing import Dict, List + +from fake_server import FakeConnection, FakeServer + +from massive import WebSocketClient +from massive.exceptions import AuthError +from massive.websocket import EquityTrade + + +def trade(sym: str, price: float = 1.0): + return {"ev": "T", "sym": sym, "x": 4, "p": price, "s": 100, "t": 1, "q": 1} + + +class FakeServerTest(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self) -> None: + self.server = await FakeServer().start() + self.addAsyncCleanup(self.server.stop) + + def client(self, **kwargs) -> WebSocketClient: + kwargs.setdefault("api_key", "key") + return WebSocketClient(feed=self.server.feed, secure=False, **kwargs) + + async def test_connect_auth_subscribe_and_receive(self) -> None: + c = self.client(subscriptions=["T.AAPL"]) + received: List[EquityTrade] = [] + + async def processor(msgs): + received.extend(msgs) + + task = asyncio.create_task(c.connect(processor)) + await self.server.wait_for_subscriptions({"T.AAPL"}) + conn = self.server.connection + self.assertEqual(conn.path, "/stocks") + self.assertEqual(conn.api_key, "key") + self.assertEqual( + conn.inbox, + [ + {"action": "auth", "params": "key"}, + {"action": "subscribe", "params": "T.AAPL"}, + ], + ) + + await self.server.send_data([trade("AAPL", 2.5)]) + await self.server.wait_for(lambda: len(received) == 1) + self.assertEqual(received[0].symbol, "AAPL") + self.assertEqual(received[0].price, 2.5) + + c.subscribe("T.MSFT") + c.unsubscribe("T.AAPL") + await self.server.wait_for_subscriptions({"T.MSFT"}) + self.assertEqual(c.subs, {"T.MSFT"}) + + await c.close() + await asyncio.wait_for(task, 5) + await self.server.wait_for(lambda: not conn.open) + + async def test_auth_failed_raises(self) -> None: + self.server.reject_keys = {"bad"} + c = self.client(api_key="bad") + with self.assertRaises(AuthError): + await asyncio.wait_for(c.connect(lambda msgs: asyncio.sleep(0)), 5) + self.assertEqual(self.server.inbox, [{"action": "auth", "params": "bad"}]) + + async def test_resubscribes_after_reconnect(self) -> None: + c = self.client(subscriptions=["T.AAPL", "Q.AAPL"]) + task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0))) + await self.server.wait_for_subscriptions({"T.AAPL", "Q.AAPL"}) + + # an abnormal close makes the client reconnect and subscribe again + await self.server.close_connections(code=1011) + await self.server.wait_for_connections(2) + await self.server.wait_for_subscriptions({"T.AAPL", "Q.AAPL"}) + self.assertEqual(len(self.server.open_connections), 1) + + await c.close() + await asyncio.wait_for(task, 5) + + async def test_scripted_auth_reply_and_messages(self) -> None: + # the fake server can replace its replies, e.g. to emulate other statuses + replies: List[str] = [] + + async def on_message(conn: FakeConnection, msg): + if msg.get("action") == "subscribe": + await conn.send([trade(t.split(".", 1)[1]) for t in conn.subscriptions]) + + self.server.on_message = on_message + self.server.ack_subscriptions = False + self.server.script_auth( + {"ev": "status", "status": "auth_success", "message": "scripted"} + ) + c = self.client(subscriptions=["T.AMZN"]) + + async def processor(msgs): + replies.extend(m.symbol for m in msgs) + await c.close() + + await asyncio.wait_for(c.connect(processor), 5) + self.assertEqual(replies, ["AMZN"]) + self.assertTrue(self.server.connection.authenticated) + + async def test_rejected_handshake_is_retried(self) -> None: + self.server.reject_handshakes = True + c = self.client(subscriptions=["T.AAPL"]) + task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0))) + await asyncio.sleep(0.2) + self.assertEqual(self.server.connections, []) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + + +class ConcurrentClientsTest(unittest.IsolatedAsyncioTestCase): + """Independent clients, each with its own thread and event loop.""" + + threads = 16 + + async def test_clients_in_parallel_threads(self) -> None: + async def on_subscribe(conn: FakeConnection, msg): + if msg.get("action") == "subscribe": + await conn.send( + [trade(t.split(".", 1)[1]) for t in msg["params"].split(",")] + ) + + server = await FakeServer(on_message=on_subscribe).start() + self.addAsyncCleanup(server.stop) + barrier = threading.Barrier(self.threads, timeout=10) + received: Dict[int, List[str]] = {} + + def run(i: int): + c = WebSocketClient( + api_key=f"key{i}", + feed=server.feed, + secure=False, + subscriptions=[f"T.S{i}"], + ) + + async def processor(msgs): + received[i] = [m.symbol for m in msgs] + await c.close() + + barrier.wait() + asyncio.run(asyncio.wait_for(c.connect(processor), 10)) + + loop = asyncio.get_running_loop() + with ThreadPoolExecutor(self.threads) as pool: + await asyncio.gather( + *(loop.run_in_executor(pool, run, i) for i in range(self.threads)) + ) + + self.assertEqual(received, {i: [f"S{i}"] for i in range(self.threads)}) + self.assertEqual(len(server.connections), self.threads) + by_key = {c.api_key: c for c in server.connections} + for i in range(self.threads): + self.assertEqual( + by_key[f"key{i}"].inbox[1], {"action": "subscribe", "params": f"T.S{i}"} + ) + + +if __name__ == "__main__": + unittest.main() From 1e510794e9fd3b8a901d13aec6c9e222c170f21d Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 15:21:20 -0600 Subject: [PATCH 2/9] test: stress a shared RESTClient from many threads 16 threads make 200 requests each through one client against a local HTTP server and check every response matches its request and per-request options don't leak into the client's default headers. Co-Authored-By: Claude Opus 5.5 --- test_rest/test_concurrency.py | 118 ++++++++++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 test_rest/test_concurrency.py diff --git a/test_rest/test_concurrency.py b/test_rest/test_concurrency.py new file mode 100644 index 00000000..1b1c4602 --- /dev/null +++ b/test_rest/test_concurrency.py @@ -0,0 +1,118 @@ +import json +import threading +import unittest +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import List, Tuple +from urllib.parse import parse_qsl, urlparse + +import pook # type: ignore + +from massive import RESTClient +from massive.rest.models.request import RequestOptionBuilder, X_MASSIVE_EDGE_ID + + +class EchoHandler(BaseHTTPRequestHandler): + # keep-alive, so the client reuses pooled connections across threads + protocol_version = "HTTP/1.1" + + def do_GET(self): + url = urlparse(self.path) + body = json.dumps( + { + "results": { + "path": url.path, + "query": dict(parse_qsl(url.query)), + "edge_id": self.headers.get(X_MASSIVE_EDGE_ID), + "authorization": self.headers.get("Authorization"), + } + } + ).encode() + with self.server.lock: # type: ignore[attr-defined] + self.server.count += 1 # type: ignore[attr-defined] + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format, *args): + pass + + +class ConcurrencyTest(unittest.TestCase): + """One RESTClient shared by many threads against a local HTTP server.""" + + threads = 16 + calls = 200 + + @classmethod + def setUpClass(cls): + # other suites leave pook intercepting urllib3; let requests through + cls.pook_was_active = pook.isactive() + pook.disable() + + @classmethod + def tearDownClass(cls): + if cls.pook_was_active: + pook.activate() + + def setUp(self): + httpd = ThreadingHTTPServer(("127.0.0.1", 0), EchoHandler) + httpd.daemon_threads = True + httpd.lock = threading.Lock() # type: ignore[attr-defined] + httpd.count = 0 # type: ignore[attr-defined] + thread = threading.Thread(target=httpd.serve_forever, daemon=True) + thread.start() + self.addCleanup(thread.join) + self.addCleanup(httpd.server_close) + self.addCleanup(httpd.shutdown) + self.httpd = httpd + host, port = httpd.server_address[:2] + self.base = f"http://{host}:{port}" + + def test_shared_client(self): + c = RESTClient( + "key", base=self.base, retries=0, maxsize=self.threads, pagination=False + ) + headers_before = dict(c.headers) + barrier = threading.Barrier(self.threads, timeout=10) + + def run(t: int) -> List[Tuple[dict, dict]]: + barrier.wait() + results = [] + for i in range(self.calls): + expected = { + "path": f"/v1/thread/{t}/call/{i}", + "query": {"t": str(t), "i": str(i)}, + "edge_id": f"edge-{t}-{i}" if i % 2 else None, + "authorization": "Bearer key", + } + options = ( + RequestOptionBuilder(edge_id=f"edge-{t}-{i}", edge_ip_address="ip") + if i % 2 + else None + ) + got = c._get( + path=expected["path"], + params={"t": t, "i": i}, + result_key="results", + options=options, + ) + results.append((expected, got)) + return results + + with ThreadPoolExecutor(self.threads) as pool: + futures = [pool.submit(run, t) for t in range(self.threads)] + results = [r for f in futures for r in f.result()] + + self.assertEqual(len(results), self.threads * self.calls) + for expected, got in results: + self.assertEqual(got, expected) + self.assertEqual(self.httpd.count, self.threads * self.calls) # type: ignore[attr-defined] + # per-request options must not leak into the shared default headers + self.assertEqual(c.headers, headers_before) + + +if __name__ == "__main__": + unittest.main() From 2c4cf6307d41fb6cd9a20336134c7e1bf2ac014a Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 20:55:09 -0600 Subject: [PATCH 3/9] build: move dev dependencies to Poetry dependency groups Replace the deprecated [tool.poetry.dev-dependencies] section with a `test` group (pook, all the unit tests need) and a `dev` group (linters, type stubs, Sphinx, orjson for the examples). `poetry install` still installs both, so local and CI workflows are unchanged, while `poetry install --only main,test` installs just what `make test` needs. Versions are unchanged. Re-locking also merged a stale duplicate entry for typing-extensions (4.12.2 below Python 3.15, 4.16.0 from 3.15) into 4.16.0, a transitive dependency of the dev tools. Co-Authored-By: Claude Opus 5.5 --- poetry.lock | 34 ++++++++++------------------------ pyproject.toml | 6 ++++-- 2 files changed, 14 insertions(+), 26 deletions(-) diff --git a/poetry.lock b/poetry.lock index 64a7a03e..1c5e5a7b 100644 --- a/poetry.lock +++ b/poetry.lock @@ -85,7 +85,7 @@ version = "22.1.0" description = "Classes Without Boilerplate" optional = false python-versions = ">=3.5" -groups = ["dev"] +groups = ["test"] files = [ {file = "attrs-22.1.0-py2.py3-none-any.whl", hash = "sha256:86efa402f67bf2df34f51a335487cf46b1ec130d02b8d39fd248abfd30da551c"}, {file = "attrs-22.1.0.tar.gz", hash = "sha256:29adc2665447e5191d0e7c568fde78b21f9672d344281d0c6e1ab085429b22b6"}, @@ -214,7 +214,7 @@ description = "Cross-platform colored terminal text." optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7" groups = ["dev"] -markers = "sys_platform == \"win32\" or platform_system == \"Windows\"" +markers = "platform_system == \"Windows\" or sys_platform == \"win32\"" files = [ {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, @@ -238,7 +238,7 @@ version = "2.1.3" description = "URL manipulation made simple." optional = false python-versions = "*" -groups = ["dev"] +groups = ["test"] files = [ {file = "furl-2.1.3-py2.py3-none-any.whl", hash = "sha256:9ab425062c4217f9802508e45feb4a83e54324273ac4b202f1850363309666c0"}, {file = "furl-2.1.3.tar.gz", hash = "sha256:5a6188fe2666c484a12159c18be97a1977a71d632ef5bb867ef15f54af39cc4e"}, @@ -299,7 +299,7 @@ version = "4.17.1" description = "An implementation of JSON Schema validation for Python" optional = false python-versions = ">=3.7" -groups = ["dev"] +groups = ["test"] files = [ {file = "jsonschema-4.17.1-py3-none-any.whl", hash = "sha256:410ef23dcdbca4eaedc08b850079179883c2ed09378bd1f760d4af4aacfa28d7"}, {file = "jsonschema-4.17.1.tar.gz", hash = "sha256:05b2d22c83640cde0b7e0aa329ca7754fbd98ea66ad8ae24aa61328dfe057fa3"}, @@ -609,7 +609,7 @@ version = "1.0.1" description = "Ordered Multivalue Dictionary" optional = false python-versions = "*" -groups = ["dev"] +groups = ["test"] files = [ {file = "orderedmultidict-1.0.1-py2.py3-none-any.whl", hash = "sha256:43c839a17ee3cdd62234c47deca1a8508a3f2ca1d0678a3bf791c87cf84adbf3"}, {file = "orderedmultidict-1.0.1.tar.gz", hash = "sha256:04070bbb5e87291cc9bfa51df413677faf2141c73c61d2a5f7b26bea3cd882ad"}, @@ -738,7 +738,7 @@ version = "2.1.6" description = "HTTP traffic mocking and expectations made easy" optional = false python-versions = ">=3.10" -groups = ["dev"] +groups = ["test"] files = [ {file = "pook-2.1.6-py3-none-any.whl", hash = "sha256:e14e710118a9193cb0e08bce2908309b9e0aa1b213b4a151c235c1f549f6eb55"}, {file = "pook-2.1.6.tar.gz", hash = "sha256:7cfcc98ffd1452d7174d0025fef7949549d8f2604a6f11de4d738647fa0ae9ec"}, @@ -770,7 +770,7 @@ version = "0.19.2" description = "Persistent/Functional/Immutable data structures" optional = false python-versions = ">=3.7" -groups = ["dev"] +groups = ["test"] files = [ {file = "pyrsistent-0.19.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:d6982b5a0237e1b7d876b60265564648a69b14017f3b5f908c5be2de3f9abb7a"}, {file = "pyrsistent-0.19.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:187d5730b0507d9285a96fca9716310d572e5464cadd19f22b63a6976254d77a"}, @@ -879,7 +879,7 @@ version = "1.16.0" description = "Python 2 and 3 compatibility utilities" optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*" -groups = ["dev"] +groups = ["test"] files = [ {file = "six-1.16.0-py2.py3-none-any.whl", hash = "sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254"}, {file = "six-1.16.0.tar.gz", hash = "sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926"}, @@ -1137,19 +1137,6 @@ files = [ {file = "types_urllib3-1.26.25.14-py3-none-any.whl", hash = "sha256:9683bbb7fb72e32bfe9d2be6e04875fbe1b3eeec3cbb4ea231435aa7fd6b4f0e"}, ] -[[package]] -name = "typing-extensions" -version = "4.12.2" -description = "Backported and Experimental Type Hints for Python 3.8+" -optional = false -python-versions = ">=3.8" -groups = ["dev"] -markers = "python_version < \"3.15\"" -files = [ - {file = "typing_extensions-4.12.2-py3-none-any.whl", hash = "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d"}, - {file = "typing_extensions-4.12.2.tar.gz", hash = "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8"}, -] - [[package]] name = "typing-extensions" version = "4.16.0" @@ -1157,7 +1144,6 @@ description = "Backported and Experimental Type Hints for Python 3.9+" optional = false python-versions = ">=3.9" groups = ["dev"] -markers = "python_version >= \"3.15\"" files = [ {file = "typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8"}, {file = "typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5"}, @@ -1266,7 +1252,7 @@ version = "0.13.0" description = "Makes working with XML feel like you are working with JSON" optional = false python-versions = ">=3.4" -groups = ["dev"] +groups = ["test"] files = [ {file = "xmltodict-0.13.0-py2.py3-none-any.whl", hash = "sha256:aa89e8fd76320154a40d19a0df04a4695fb9dc5ba977cbb68ab3e4eb225e7852"}, {file = "xmltodict-0.13.0.tar.gz", hash = "sha256:341595a488e3e01a85a9d8911d8912fd922ede5fecc4dce437eb4b6c8d037e56"}, @@ -1275,4 +1261,4 @@ files = [ [metadata] lock-version = "2.1" python-versions = "^3.10" -content-hash = "610b6cf2d6f3cc2100ae3173b2af859dc7e399e8263d67c6586db113e766ca3c" +content-hash = "50ca3ffb193a876b67e29ff447d8bbab038cd5487e0b6bafb2b0cbe7b14aa2e0" diff --git a/pyproject.toml b/pyproject.toml index bb58f3da..4483e735 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,10 @@ urllib3 = ">=1.26.9" websockets = ">=14.0" certifi = ">=2022.5.18,<2027.0.0" -[tool.poetry.dev-dependencies] +[tool.poetry.group.test.dependencies] +pook = "^2.1.6" + +[tool.poetry.group.dev.dependencies] black = "^26.3.1" mypy = "^2.4" types-urllib3 = "^1.26.25" @@ -39,7 +42,6 @@ sphinx-rtd-theme = "^3.1.0" sphinx-autodoc-typehints = "^2.3.0" types-certifi = "^2021.10.8" types-setuptools = "^84.0.0" -pook = "^2.1.6" orjson = "^3.11.6" [build-system] From af964947bf2e40204a82c93c5bf3109f56b46a97 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 20:55:16 -0600 Subject: [PATCH 4/9] ci: run unit tests on free-threaded Python 3.13t and 3.14t Add a job that runs `make test` on 3.13t and 3.14t with PYTHON_GIL=0, and extend the regular matrix to 3.12, 3.13 and 3.14. The free-threaded job installs only the main and test dependency groups: the dev group is not needed to run the tests, and orjson in it does not support free-threaded Python. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/test.yml | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index a1ae9ada..e49a8035 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -14,7 +14,7 @@ jobs: fail-fast: false matrix: os: [ubuntu-latest, windows-latest] - python-version: ['3.10', '3.11'] + python-version: ['3.10', '3.11', '3.12', '3.13', '3.14'] runs-on: ${{ matrix.os }} name: ${{ matrix.os }} Unit test ${{ matrix.python-version }} steps: @@ -32,3 +32,32 @@ jobs: run: poetry install - name: Unit tests run: make test + test-free-threaded: + strategy: + fail-fast: false + matrix: + python-version: ['3.13t', '3.14t'] + runs-on: ubuntu-latest + name: ubuntu-latest Unit test ${{ matrix.python-version }} (PYTHON_GIL=0) + steps: + - uses: actions/checkout@v3 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Setup Poetry + # pipx is Poetry's recommended installation method, and it keeps Poetry + # on the runner's regular Python, independent of the interpreter under test + run: pipx install poetry + - name: Install pypi deps + # the dev group (linters, docs, orjson for the examples) is not needed + # to run the tests, and orjson does not support free-threaded Python + run: | + poetry env use python${{ matrix.python-version }} + poetry install --only main,test + - name: Unit tests + env: + PYTHON_GIL: '0' + run: | + poetry run python -c 'import sys; assert not sys._is_gil_enabled(), sys.version' + make test From 84931d248d45573cc445457dada0131679303290 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 15:40:44 -0600 Subject: [PATCH 5/9] fix!: reconnect after the server closes the connection cleanly connect() returned on ConnectionClosedOK, so a clean close from the server, which it sends e.g. after a slow_consumer or max_connections status, ended the client for good. Treat it like any other disconnect and reconnect, unless close() was called. On any disconnect, keep scheduled_subs, the desired set: it used to be replaced with the topics sent so far, so a subscribe() or unsubscribe() made just before the drop was lost. Reset only subs, so the next connection subscribes exactly the desired set. When connect() gives up after max_reconnects, it raises the error that ended the last connection, ConnectionClosedOK included. The status test that relied on connect() returning now sets max_reconnects=0. BREAKING CHANGE: connect() no longer returns when the server closes the connection cleanly; it reconnects until close() is called or max_reconnects is exceeded. --- massive/websocket/__init__.py | 51 +++++++++++++++++---- test_websocket/test_reconnect.py | 79 ++++++++++++++++++++++++++++++++ test_websocket/test_status.py | 11 +++-- 3 files changed, 129 insertions(+), 12 deletions(-) create mode 100644 test_websocket/test_reconnect.py diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 7463b4cb..703244bf 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -6,10 +6,14 @@ import json import asyncio import ssl +import threading import time from .models import * from websockets.asyncio.client import connect, ClientConnection -from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError +from websockets.exceptions import ( + ConnectionClosed, + ConnectionClosedOK, +) from ..certs import ca_bundle_path from ..logging import get_logger import logging @@ -85,6 +89,9 @@ def __init__( self.json = custom_json else: self.json = json + # set while connect() runs; close() also stops it from reconnecting + self._loop: Optional[asyncio.AbstractEventLoop] = None + self._close_requested = threading.Event() # https://websockets.readthedocs.io/en/stable/reference/client.html#opening-a-connection async def connect( @@ -99,10 +106,33 @@ async def connect( """ Connect to websocket server and run `processor(msg)` on every new `msg`. + When the connection ends, whether the server closed it cleanly or not, + the client reconnects and subscribes again to every scheduled + subscription. It returns once `close()` is called. + :param processor: The callback to process messages. :param close_timeout: How long to wait for handshake when calling .close. :raises AuthError: If invalid API key is supplied. + :raises ConnectionClosed: If more than `max_reconnects` reconnects fail, + the error that ended the last connection; ConnectionClosedOK if the + server closed it cleanly. """ + self._loop = asyncio.get_running_loop() + try: + await self._connect(processor, close_timeout, **kwargs) + finally: + self._loop = None + self._close_requested.clear() + + async def _connect( + self, + processor: Union[ + Callable[[List[WebSocketMessage]], Awaitable], + Callable[[Union[str, bytes]], Awaitable], + ], + close_timeout: int, + **kwargs, + ): reconnects = 0 logger.debug("connect: %s", self.url) ssl_context = self._ssl_context() @@ -165,17 +195,18 @@ async def connect( if len(cmsg) > 0: await processor(cmsg) # type: ignore - except ConnectionClosedOK as e: - logger.debug("connection closed (OK): %s", e) + except ConnectionClosed as e: + ok = isinstance(e, ConnectionClosedOK) + logger.debug("connection closed (%s): %s", "OK" if ok else "ERR", e) # the server sends a status (e.g. slow_consumer) before closing self._flush_server_errors(force=True) - return - except ConnectionClosedError as e: + if self._close_requested.is_set(): + return + # the server can close cleanly too, e.g. after max_connections last_exc = e - logger.debug("connection closed (ERR): %s", e) - self._flush_server_errors(force=True) reconnects += 1 - self.scheduled_subs = set(self.subs) + # scheduled_subs is the desired set, including changes not sent + # yet: keep it, and subscribe all of it on the next connection self.subs = set() self.schedule_resub = True if self.max_reconnects is not None and reconnects > self.max_reconnects: @@ -320,9 +351,11 @@ def unsubscribe_all(self): async def close(self): """ - Close the websocket connection. + Close the websocket connection and stop `connect()` from reconnecting. """ logger.debug("closing") + if self._loop is not None: + self._close_requested.set() self._flush_server_errors(force=True) if self.websocket: diff --git a/test_websocket/test_reconnect.py b/test_websocket/test_reconnect.py new file mode 100644 index 00000000..c98646fc --- /dev/null +++ b/test_websocket/test_reconnect.py @@ -0,0 +1,79 @@ +import asyncio +import unittest +from typing import Any, List + +from fake_server import FakeServer + +from massive import WebSocketClient + + +async def ignore(msgs: Any) -> None: + pass + + +class ReconnectTest(unittest.IsolatedAsyncioTestCase): + """How the client behaves when its connection ends.""" + + async def asyncSetUp(self) -> None: + self.server = await FakeServer().start() + self.addAsyncCleanup(self.server.stop) + + def client(self, **kwargs) -> WebSocketClient: + kwargs.setdefault("api_key", "key") + return WebSocketClient(feed=self.server.feed, secure=False, **kwargs) + + def start(self, c: WebSocketClient, processor=ignore) -> "asyncio.Task[None]": + task = asyncio.create_task(c.connect(processor)) + + async def cancel(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + self.addAsyncCleanup(cancel) + return task + + async def test_keeps_changes_scheduled_during_an_outage(self) -> None: + c = self.client(subscriptions=["T.KEEP", "T.OLD", "Q.OLD"]) + task = self.start(c) + await self.server.wait_for_subscriptions({"T.KEEP", "T.OLD", "Q.OLD"}) + + self.server.reject_handshakes = True + # scheduled but not sent yet when the connection drops + c.subscribe("T.NEW") + c.unsubscribe("T.OLD") + await self.server.drop_connections() + # scheduled while the client cannot reconnect + c.subscribe("Q.NEW") + c.unsubscribe("Q.OLD") + self.server.reject_handshakes = False + + await self.server.wait_for_connections(2, timeout=15) + await self.server.wait_for_subscriptions({"T.KEEP", "T.NEW", "Q.NEW"}) + topics: List[str] = [] + for m in self.server.connection.inbox[1:]: + self.assertEqual(m["action"], "subscribe") + topics += m["params"].split(",") + self.assertCountEqual(topics, ["T.KEEP", "T.NEW", "Q.NEW"]) + self.assertEqual(c.scheduled_subs, {"T.KEEP", "T.NEW", "Q.NEW"}) + + await c.close() + await asyncio.wait_for(task, 5) + + async def test_reconnects_after_a_clean_close_until_closed(self) -> None: + c = self.client(subscriptions=["T.AAPL"]) + task = self.start(c) + await self.server.wait_for_subscriptions({"T.AAPL"}) + + # e.g. the server closes cleanly after a slow_consumer status + await self.server.close_connections(code=1000) + await self.server.wait_for_connections(2) + await self.server.wait_for_subscriptions({"T.AAPL"}) + self.assertFalse(task.done()) + + await c.close() + await asyncio.wait_for(task, 5) + self.assertEqual(len(self.server.connections), 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test_websocket/test_status.py b/test_websocket/test_status.py index d333c9c8..1212e38e 100644 --- a/test_websocket/test_status.py +++ b/test_websocket/test_status.py @@ -3,6 +3,7 @@ from base_ws import BaseTest from mock_server import port from websockets.asyncio.client import connect +from websockets.exceptions import ConnectionClosedOK import asyncio @@ -65,11 +66,15 @@ async def test_success_statuses_are_not_logged(self): self.assertEqual(logs.output, ["WARNING:WebSocketClient:sentinel"]) async def test_disconnect_reason_is_logged(self): - # the server sends a status and then closes the connection normally + # the server sends a status and then closes the connection normally; + # the client would reconnect, so let it give up instead with self.assertLogs("WebSocketClient", level="WARNING") as logs: - c = WebSocketClient(api_key="", feed=f"localhost:{port}", secure=False) + c = WebSocketClient( + api_key="", feed=f"localhost:{port}", secure=False, max_reconnects=0 + ) c.subscribe("SLOW.X") - await asyncio.wait_for(c.connect(lambda msgs: asyncio.sleep(0)), 5) + with self.assertRaises(ConnectionClosedOK): + await asyncio.wait_for(c.connect(lambda msgs: asyncio.sleep(0)), 5) self.assertEqual( logs.output, [ From 0de5fe1d77d60d3993bcfd2e7e3f15aeeb744edb Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:23:28 -0600 Subject: [PATCH 6/9] fix!: wait with a doubling, jittered backoff before reconnecting Reconnects were immediate after a disconnect: the websockets connect() iterator resets its delay after every successful connection, and its own delay for refused handshakes starts below 5 s. Massive asks clients to wait at least 15 s. The iterator also retried every error, including ones that cannot pass, such as HTTP 401 or an unverifiable certificate. connect() now makes each attempt itself, with a new websockets connection object per attempt; one is built up front so that invalid arguments still raise at once. Only transient failures are retried: OSError (refused, unreachable, timed out, a DNS failure, or a TLS error other than certificate verification), the connection closing before a response (EOFError in websockets 14, InvalidMessage caused by EOFError or OSError in websockets 15), and HTTP 408, 429 and 5xx. Anything else is raised at once. A process_exception argument, which only the iterator used, is rejected with TypeError. After the n-th consecutive failure the client logs a warning and waits uniform(b, min(1.2 * b, max_reconnect_backoff)) seconds, where b = min(reconnect_backoff * 2**n, max_reconnect_backoff). Jitter only lengthens a wait and every wait lies in [reconnect_backoff, max_reconnect_backoff]. The count resets once a connection authenticates. close() during the wait returns at once, without the "no websocket open to close" warning. New constructor parameters: reconnect_backoff (default 15.0) and max_reconnect_backoff (default 120.0). BREAKING CHANGE: reconnects wait at least 15 s by default (they used to be immediate); pass reconnect_backoff to change it. Errors that cannot pass, e.g. HTTP 4xx other than 408 and 429, are raised instead of retried. --- massive/websocket/__init__.py | 134 +++++++++++++++-- test_websocket/test_fake_server.py | 17 ++- test_websocket/test_reconnect.py | 232 ++++++++++++++++++++++++++++- 3 files changed, 359 insertions(+), 24 deletions(-) diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 703244bf..b9e2103e 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -5,6 +5,7 @@ import logging import json import asyncio +import random import ssl import threading import time @@ -13,6 +14,8 @@ from websockets.exceptions import ( ConnectionClosed, ConnectionClosedOK, + InvalidMessage, + InvalidStatus, ) from ..certs import ca_bundle_path from ..logging import get_logger @@ -40,6 +43,8 @@ def __init__( secure: bool = True, custom_json: Optional[Any] = None, ca_certs: Optional[str] = None, + reconnect_backoff: float = 15.0, + max_reconnect_backoff: float = 120.0, **kwargs, ): """ @@ -53,8 +58,22 @@ def __init__( :param max_reconnects: How many times to reconnect on network outage before ending .connect event loop. :param custom_json: Optional module exposing loads/dumps functions (similar to Python's json module) to be used for JSON conversions. :param ca_certs: Path to a CA bundle to verify TLS with. Defaults to the SSL_CERT_FILE or REQUESTS_CA_BUNDLE environment variable, then certifi. + :param reconnect_backoff: Seconds to wait before reconnecting after the + connection ends or an attempt to connect fails. The wait doubles with + each consecutive failure and gets up to 20% random jitter, but never + leaves [reconnect_backoff, max_reconnect_backoff]. Defaults to 15, the + minimum Massive asks clients to wait; 0 disables the wait. + :param max_reconnect_backoff: The longest wait between reconnect + attempts, jitter included. + :raises ValueError: If reconnect_backoff is negative or greater than + max_reconnect_backoff. :return: A client. """ + if reconnect_backoff < 0 or max_reconnect_backoff < reconnect_backoff: + raise ValueError( + "need 0 <= reconnect_backoff <= max_reconnect_backoff, got " + f"{reconnect_backoff} and {max_reconnect_backoff}" + ) if api_key is None: raise AuthError( f"Must specify env var {env_key} or pass api_key in constructor" @@ -92,6 +111,9 @@ def __init__( # set while connect() runs; close() also stops it from reconnecting self._loop: Optional[asyncio.AbstractEventLoop] = None self._close_requested = threading.Event() + self._close_event: Optional[asyncio.Event] = None + self.reconnect_backoff = reconnect_backoff + self.max_reconnect_backoff = max_reconnect_backoff # https://websockets.readthedocs.io/en/stable/reference/client.html#opening-a-connection async def connect( @@ -107,21 +129,32 @@ async def connect( Connect to websocket server and run `processor(msg)` on every new `msg`. When the connection ends, whether the server closed it cleanly or not, - the client reconnects and subscribes again to every scheduled - subscription. It returns once `close()` is called. + or an attempt to connect fails for a reason that may pass, the client + waits (see `reconnect_backoff`), reconnects, and subscribes again to + every scheduled subscription. Such reasons are network errors (refused, + unreachable, timed out, a DNS lookup failure, or a TLS error other + than certificate verification), the connection closing during the + handshake, and HTTP 408, 429 or 5xx. Other errors, e.g. any other HTTP + 4xx or a TLS certificate that cannot be verified, are raised at once. + It returns once `close()` is called. :param processor: The callback to process messages. :param close_timeout: How long to wait for handshake when calling .close. + :param kwargs: Passed to websockets' connect() for each attempt. :raises AuthError: If invalid API key is supplied. + :raises TypeError: If `process_exception` is passed: connect() decides + itself which errors to retry. :raises ConnectionClosed: If more than `max_reconnects` reconnects fail, the error that ended the last connection; ConnectionClosedOK if the server closed it cleanly. """ self._loop = asyncio.get_running_loop() + self._close_event = asyncio.Event() try: await self._connect(processor, close_timeout, **kwargs) finally: self._loop = None + self._close_event = None self._close_requested.clear() async def _connect( @@ -134,17 +167,41 @@ async def _connect( **kwargs, ): reconnects = 0 + # consecutive failed attempts, which set the backoff + failures = 0 logger.debug("connect: %s", self.url) ssl_context = self._ssl_context() - last_exc = None - async for s in connect( - self.url, - close_timeout=close_timeout, - ssl=ssl_context, - process_exception=lambda exc: None, - **kwargs, - ): + if "process_exception" in kwargs: + raise TypeError( + "process_exception is not supported: connect() decides which " + "errors to retry" + ) + + def new_attempt() -> Awaitable[ClientConnection]: + return connect( + self.url, close_timeout=close_timeout, ssl=ssl_context, **kwargs + ) + + # build one up front so that invalid arguments raise at once + new_attempt() + last_exc: Optional[Exception] = None + delay = 0.0 + while not self._close_requested.is_set(): + if failures > 0 and await self._wait_for_close(delay): + return + try: + s = await new_attempt() + except Exception as e: + if not self._is_transient(e): + raise + failures += 1 + delay = self._reconnect_delay(failures - 1) + logger.warning("connect failed; reconnecting in %.1f s: %s", delay, e) + continue + if self._close_requested.is_set(): + await s.close() + return self.websocket = s try: msg = await s.recv() @@ -159,6 +216,7 @@ async def _connect( logger.debug("authed: %s", auth_msg) if auth_msg_parsed[0]["status"] == "auth_failed": raise AuthError(auth_msg_parsed[0]["message"]) + failures = 0 # e.g. max_connections, which the server sends before closing self._log_server_statuses(auth_msg_parsed) self._flush_server_errors(force=True) @@ -204,14 +262,21 @@ async def _connect( return # the server can close cleanly too, e.g. after max_connections last_exc = e + failures += 1 reconnects += 1 + if self.max_reconnects is not None and reconnects > self.max_reconnects: + break + delay = self._reconnect_delay(failures - 1) + logger.warning( + "connection closed; reconnecting in %.1f s: %s", delay, e + ) + finally: + self.websocket = None # scheduled_subs is the desired set, including changes not sent # yet: keep it, and subscribe all of it on the next connection self.subs = set() self.schedule_resub = True - if self.max_reconnects is not None and reconnects > self.max_reconnects: - break - continue + await s.close() if ( last_exc @@ -220,6 +285,44 @@ async def _connect( ): raise last_exc + @staticmethod + def _is_transient(exc: Exception) -> bool: + """Whether a failed attempt to connect is worth retrying.""" + if isinstance(exc, InvalidStatus): + status = exc.response.status_code + return status >= 500 or status in (408, 429) + if isinstance(exc, InvalidMessage): + # websockets 15: the connection closed or failed before a response + return isinstance(exc.__cause__, (EOFError, OSError)) + if isinstance(exc, ssl.SSLCertVerificationError): + return False + # e.g. refused, unreachable, timed out, DNS, or another TLS error; + # websockets 14 raises EOFError if the connection closes before a response + return isinstance(exc, (OSError, EOFError, asyncio.TimeoutError)) + + def _reconnect_delay(self, attempt: int) -> float: + """ + The wait before reconnect attempt `attempt`, counting from 0. + + The base doubles from `reconnect_backoff` up to `max_reconnect_backoff`; + jitter only lengthens it, by up to 20%, and is clamped to the cap. + """ + # bound the exponent so a long outage cannot overflow the float + base = min( + self.reconnect_backoff * 2.0 ** min(attempt, 1000), + self.max_reconnect_backoff, + ) + return random.uniform(base, min(1.2 * base, self.max_reconnect_backoff)) + + async def _wait_for_close(self, timeout: float) -> bool: + """Wait up to `timeout` seconds; return whether close was requested.""" + assert self._close_event is not None + try: + await asyncio.wait_for(self._close_event.wait(), timeout) + except asyncio.TimeoutError: + pass + return self._close_requested.is_set() + def _ssl_context(self) -> Optional[ssl.SSLContext]: # darwin needs some extra <3 if not self.url.startswith("wss://"): @@ -356,10 +459,13 @@ async def close(self): logger.debug("closing") if self._loop is not None: self._close_requested.set() + if self._close_event is not None: + self._close_event.set() self._flush_server_errors(force=True) if self.websocket: await self.websocket.close() self.websocket = None - else: + elif self._loop is None: + # (while connect() runs, it may be between connections) logger.warning("no websocket open to close") diff --git a/test_websocket/test_fake_server.py b/test_websocket/test_fake_server.py index ede70ac3..9138051e 100644 --- a/test_websocket/test_fake_server.py +++ b/test_websocket/test_fake_server.py @@ -66,7 +66,7 @@ async def test_auth_failed_raises(self) -> None: self.assertEqual(self.server.inbox, [{"action": "auth", "params": "bad"}]) async def test_resubscribes_after_reconnect(self) -> None: - c = self.client(subscriptions=["T.AAPL", "Q.AAPL"]) + c = self.client(subscriptions=["T.AAPL", "Q.AAPL"], reconnect_backoff=0.01) task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0))) await self.server.wait_for_subscriptions({"T.AAPL", "Q.AAPL"}) @@ -104,13 +104,16 @@ async def processor(msgs): async def test_rejected_handshake_is_retried(self) -> None: self.server.reject_handshakes = True - c = self.client(subscriptions=["T.AAPL"]) - task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0))) - await asyncio.sleep(0.2) + c = self.client(subscriptions=["T.AAPL"], reconnect_backoff=0.01) + with self.assertLogs("WebSocketClient", level="WARNING") as logs: + task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0))) + # each refused attempt is logged before the next one + await self.server.wait_for(lambda: len(logs.records) >= 2) self.assertEqual(self.server.connections, []) - task.cancel() - with self.assertRaises(asyncio.CancelledError): - await task + self.server.reject_handshakes = False + await self.server.wait_for_subscriptions({"T.AAPL"}) + await c.close() + await asyncio.wait_for(task, 5) class ConcurrentClientsTest(unittest.IsolatedAsyncioTestCase): diff --git a/test_websocket/test_reconnect.py b/test_websocket/test_reconnect.py index c98646fc..f498d915 100644 --- a/test_websocket/test_reconnect.py +++ b/test_websocket/test_reconnect.py @@ -1,10 +1,46 @@ import asyncio +import socket +import ssl +import time import unittest -from typing import Any, List +from typing import Any, List, Optional +from unittest import mock from fake_server import FakeServer +from websockets.asyncio.server import ServerConnection +from websockets.datastructures import Headers +from websockets.exceptions import ( + InvalidMessage, + InvalidStatus, + InvalidURI, +) +from websockets.http11 import Request, Response +import massive.websocket as ws_module from massive import WebSocketClient +from massive.exceptions import AuthError + +# small waits keep the tests fast; the default is 15 s +BACKOFF = 0.01 +MAX_BACKOFF = 0.05 + + +class TimedServer(FakeServer): + """A FakeServer that records when each handshake was attempted.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.attempts: List[float] = [] + # if set, answer handshakes with this HTTP status + self.reject_status: Optional[int] = None + + def _process_request( + self, websocket: ServerConnection, request: Request + ) -> Optional[Response]: + self.attempts.append(time.monotonic()) + if self.reject_status is not None: + return websocket.respond(self.reject_status, "rejected\n") + return super()._process_request(websocket, request) async def ignore(msgs: Any) -> None: @@ -15,11 +51,14 @@ class ReconnectTest(unittest.IsolatedAsyncioTestCase): """How the client behaves when its connection ends.""" async def asyncSetUp(self) -> None: - self.server = await FakeServer().start() + self.server = TimedServer() + await self.server.start() self.addAsyncCleanup(self.server.stop) def client(self, **kwargs) -> WebSocketClient: kwargs.setdefault("api_key", "key") + kwargs.setdefault("reconnect_backoff", BACKOFF) + kwargs.setdefault("max_reconnect_backoff", MAX_BACKOFF) return WebSocketClient(feed=self.server.feed, secure=False, **kwargs) def start(self, c: WebSocketClient, processor=ignore) -> "asyncio.Task[None]": @@ -47,7 +86,7 @@ async def test_keeps_changes_scheduled_during_an_outage(self) -> None: c.unsubscribe("Q.OLD") self.server.reject_handshakes = False - await self.server.wait_for_connections(2, timeout=15) + await self.server.wait_for_connections(2) await self.server.wait_for_subscriptions({"T.KEEP", "T.NEW", "Q.NEW"}) topics: List[str] = [] for m in self.server.connection.inbox[1:]: @@ -74,6 +113,193 @@ async def test_reconnects_after_a_clean_close_until_closed(self) -> None: await asyncio.wait_for(task, 5) self.assertEqual(len(self.server.connections), 2) + async def test_waits_between_attempts(self) -> None: + c = self.client(reconnect_backoff=0.05, max_reconnect_backoff=0.2) + delays: List[float] = [] + delay = c._reconnect_delay + + def record(attempt: int) -> float: + delays.append(delay(attempt)) + return delays[-1] + + with mock.patch.object(c, "_reconnect_delay", side_effect=record): + task = self.start(c) + await self.server.wait_for_connections(1) + self.server.reject_handshakes = True + await self.server.drop_connections() + # 0.05, 0.1, 0.2, then the cap + await self.server.wait_for(lambda: len(self.server.attempts) == 6) + + # close() during the wait stops connect() + await c.close() + await asyncio.wait_for(task, 5) + + self.assertEqual(len(self.server.attempts), 6) + # one delay before each refused attempt, maybe one more being waited + self.assertIn(len(delays), (5, 6)) + for d in delays: + self.assertGreaterEqual(d, 0.05) + self.assertLessEqual(d, 0.2) + self.assertGreaterEqual(delays[1], 0.1) + # at the cap, jitter cannot lengthen the wait + self.assertEqual(delays[2:], [0.2] * (len(delays) - 2)) + # each refused attempt waited for its delay before the next one + gaps = [ + b - a for a, b in zip(self.server.attempts[1:], self.server.attempts[2:]) + ] + # allow for the timer resolution, about 15.6 ms on Windows + for gap, d in zip(gaps, delays[1:]): + self.assertGreaterEqual(gap, d - 0.03) + + async def test_transient_handshake_errors_are_retried(self) -> None: + for code in (408, 429, 500, 503): + self.server.reject_status = code + n = len(self.server.attempts) + c = self.client() + with self.assertLogs("WebSocketClient", level="WARNING") as logs: + task = self.start(c) + await self.server.wait_for(lambda: len(self.server.attempts) >= n + 3) + self.assertIn( + f"connect failed; reconnecting in 0.0 s: server rejected WebSocket " + f"connection: HTTP {code}", + logs.output[0], + ) + self.server.reject_status = None + await self.server.wait_for( + lambda: bool(self.server.open_connections), timeout=5 + ) + await c.close() + await asyncio.wait_for(task, 5) + + async def test_refused_connections_are_retried(self) -> None: + port = self.server.port + await self.server.stop() + c = WebSocketClient( + api_key="key", feed=f"127.0.0.1:{port}", secure=False, reconnect_backoff=0 + ) + with self.assertLogs("WebSocketClient", level="WARNING") as logs: + task = self.start(c) + # Windows retries the SYN before it reports a refusal, so there an + # attempt can take about 2 s + await self.server.wait_for(lambda: len(logs.records) >= 2, timeout=15) + self.assertFalse(task.done()) + self.assertIn("connect failed; reconnecting in", logs.output[0]) + + async def test_fatal_handshake_errors_raise_at_once(self) -> None: + for code in (400, 401, 403, 404): + self.server.reject_status = code + n = len(self.server.attempts) + with self.assertRaises(InvalidStatus) as e: + await asyncio.wait_for(self.client().connect(ignore), 5) + self.assertEqual(e.exception.response.status_code, code) + self.assertEqual(len(self.server.attempts), n + 1) + + async def test_invalid_arguments_raise_at_once(self) -> None: + # rejected when building the connection + with self.assertRaises(ValueError): + await self.client(max_reconnects=None).connect(ignore, compression="x") + # rejected by the first attempt + with self.assertRaises(TypeError): + await self.client(max_reconnects=None).connect(ignore, bogus=1) + # connect() decides which errors to retry + with self.assertRaises(TypeError): + await self.client().connect(ignore, process_exception=lambda e: None) + c = WebSocketClient(api_key="key", feed="127.0.0.1:port", secure=False) + with self.assertRaises(ValueError): + await asyncio.wait_for(c.connect(ignore), 5) + self.assertEqual(self.server.attempts, []) + + +class TransientTest(unittest.TestCase): + """Which failed attempts to connect are retried.""" + + def test_classification(self) -> None: + def status(code: int) -> InvalidStatus: + return InvalidStatus(Response(code, "", Headers(), b"")) + + def invalid(cause: Optional[BaseException]) -> InvalidMessage: + # how websockets 15 reports a handshake without a valid response + e = InvalidMessage("did not receive a valid HTTP response") + e.__cause__ = cause + return e + + transient = [ + status(408), + status(429), + status(500), + status(503), + invalid(EOFError("connection closed while reading HTTP status line")), + invalid(ConnectionResetError()), + # how websockets 14 reports a connection closed before a response + EOFError("line without CRLF"), + ConnectionRefusedError(), + OSError("unreachable"), + socket.gaierror(), + TimeoutError(), + asyncio.TimeoutError(), + ] + fatal = [ + status(400), + status(401), + status(403), + status(404), + invalid(ValueError("invalid HTTP status line")), + invalid(None), + InvalidURI("ws://x", "bad"), + ssl.SSLCertVerificationError(), + TypeError(), + ValueError(), + AuthError("bad key"), + ] + for e in transient: + self.assertTrue(WebSocketClient._is_transient(e), repr(e)) + for e in fatal: + self.assertFalse(WebSocketClient._is_transient(e), repr(e)) + + +class ReconnectDelayTest(unittest.TestCase): + """The wait before each consecutive reconnect attempt.""" + + def delays(self, c: WebSocketClient, n: int) -> List[float]: + return [c._reconnect_delay(i) for i in range(n)] + + def test_defaults(self) -> None: + c = WebSocketClient(api_key="key") + self.assertEqual(c.reconnect_backoff, 15.0) + self.assertEqual(c.max_reconnect_backoff, 120.0) + with mock.patch.object(ws_module.random, "uniform", side_effect=lambda a, b: a): + self.assertEqual(self.delays(c, 6), [15, 30, 60, 120, 120, 120]) + with mock.patch.object(ws_module.random, "uniform", side_effect=lambda a, b: b): + self.assertEqual(self.delays(c, 6), [18, 36, 72, 120, 120, 120]) + + def test_jitter_is_clamped_to_the_cap(self) -> None: + c = WebSocketClient( + api_key="key", reconnect_backoff=10, max_reconnect_backoff=100 + ) + with mock.patch.object(ws_module.random, "uniform", side_effect=lambda a, b: a): + self.assertEqual(self.delays(c, 6), [10, 20, 40, 80, 100, 100]) + with mock.patch.object(ws_module.random, "uniform", side_effect=lambda a, b: b): + self.assertEqual(self.delays(c, 6), [12, 24, 48, 96, 100, 100]) + + def test_every_delay_is_within_bounds(self) -> None: + c = WebSocketClient(api_key="key") + for attempt in list(range(20)) * 50 + [10**6]: + d = c._reconnect_delay(attempt) + self.assertGreaterEqual(d, 15.0) + self.assertLessEqual(d, 120.0) + + def test_equal_bounds(self) -> None: + c = WebSocketClient(api_key="key", reconnect_backoff=0, max_reconnect_backoff=0) + self.assertEqual(self.delays(c, 3), [0, 0, 0]) + + def test_invalid_bounds(self) -> None: + with self.assertRaises(ValueError): + WebSocketClient(api_key="key", reconnect_backoff=-1) + with self.assertRaises(ValueError): + WebSocketClient( + api_key="key", reconnect_backoff=30, max_reconnect_backoff=20 + ) + if __name__ == "__main__": unittest.main() From 3bd7c3eafcb75248a276b8a3a52cf06d90b46544 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:23:55 -0600 Subject: [PATCH 7/9] fix!: count consecutive failures toward max_reconnects max_reconnects limited the reconnects over the client's lifetime, so with the default of 5 a long-running client raised on its sixth outage even if every reconnect had succeeded. Count consecutive failed attempts instead (refused connections included) and start over whenever a connection authenticates. The default becomes None, retrying forever. Only transient failures are retried, with a growing wait, and other errors are raised at once, so run() is safe in a long-lived process. BREAKING CHANGE: max_reconnects now defaults to None (retry forever) instead of 5, and counts consecutive failed attempts rather than all reconnects. --- massive/websocket/__init__.py | 35 +++++++++++++++----------- test_websocket/test_reconnect.py | 42 +++++++++++++++++++++++++++++++- 2 files changed, 62 insertions(+), 15 deletions(-) diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index b9e2103e..a651eb99 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -39,7 +39,7 @@ def __init__( raw: bool = False, verbose: bool = False, subscriptions: Optional[List[str]] = None, - max_reconnects: Optional[int] = 5, + max_reconnects: Optional[int] = None, secure: bool = True, custom_json: Optional[Any] = None, ca_certs: Optional[str] = None, @@ -55,7 +55,13 @@ def __init__( :param raw: Whether to pass raw Union[str, bytes] to user callback. :param verbose: Whether to log client and server status messages. :param subscriptions: List of subscription parameters. - :param max_reconnects: How many times to reconnect on network outage before ending .connect event loop. + :param max_reconnects: How many consecutive reconnect attempts may fail + (the connection is refused, or ends before or after authenticating) + before .connect raises the last error. The count starts over whenever + a connection authenticates. None, the default, retries forever. Only + transient failures are retried and the wait between attempts grows, + while other errors (e.g. HTTP 4xx) are raised at once, so with the + default .run is safe in a long-lived process. :param custom_json: Optional module exposing loads/dumps functions (similar to Python's json module) to be used for JSON conversions. :param ca_certs: Path to a CA bundle to verify TLS with. Defaults to the SSL_CERT_FILE or REQUESTS_CA_BUNDLE environment variable, then certifi. :param reconnect_backoff: Seconds to wait before reconnecting after the @@ -144,9 +150,9 @@ async def connect( :raises AuthError: If invalid API key is supplied. :raises TypeError: If `process_exception` is passed: connect() decides itself which errors to retry. - :raises ConnectionClosed: If more than `max_reconnects` reconnects fail, - the error that ended the last connection; ConnectionClosedOK if the - server closed it cleanly. + :raises Exception: If more than `max_reconnects` consecutive attempts + fail, the error that ended the last one, e.g. ConnectionClosedOK if + the server closed the last connection cleanly. """ self._loop = asyncio.get_running_loop() self._close_event = asyncio.Event() @@ -166,8 +172,8 @@ async def _connect( close_timeout: int, **kwargs, ): - reconnects = 0 - # consecutive failed attempts, which set the backoff + # consecutive failed attempts, which set the backoff and are limited + # by max_reconnects failures = 0 logger.debug("connect: %s", self.url) ssl_context = self._ssl_context() @@ -195,7 +201,10 @@ def new_attempt() -> Awaitable[ClientConnection]: except Exception as e: if not self._is_transient(e): raise + last_exc = e failures += 1 + if self._too_many_failures(failures): + break delay = self._reconnect_delay(failures - 1) logger.warning("connect failed; reconnecting in %.1f s: %s", delay, e) continue @@ -263,8 +272,7 @@ def new_attempt() -> Awaitable[ClientConnection]: # the server can close cleanly too, e.g. after max_connections last_exc = e failures += 1 - reconnects += 1 - if self.max_reconnects is not None and reconnects > self.max_reconnects: + if self._too_many_failures(failures): break delay = self._reconnect_delay(failures - 1) logger.warning( @@ -278,13 +286,12 @@ def new_attempt() -> Awaitable[ClientConnection]: self.schedule_resub = True await s.close() - if ( - last_exc - and self.max_reconnects is not None - and reconnects > self.max_reconnects - ): + if last_exc is not None and self._too_many_failures(failures): raise last_exc + def _too_many_failures(self, failures: int) -> bool: + return self.max_reconnects is not None and failures > self.max_reconnects + @staticmethod def _is_transient(exc: Exception) -> bool: """Whether a failed attempt to connect is worth retrying.""" diff --git a/test_websocket/test_reconnect.py b/test_websocket/test_reconnect.py index f498d915..33343403 100644 --- a/test_websocket/test_reconnect.py +++ b/test_websocket/test_reconnect.py @@ -114,7 +114,9 @@ async def test_reconnects_after_a_clean_close_until_closed(self) -> None: self.assertEqual(len(self.server.connections), 2) async def test_waits_between_attempts(self) -> None: - c = self.client(reconnect_backoff=0.05, max_reconnect_backoff=0.2) + c = self.client( + reconnect_backoff=0.05, max_reconnect_backoff=0.2, max_reconnects=None + ) delays: List[float] = [] delay = c._reconnect_delay @@ -209,6 +211,43 @@ async def test_invalid_arguments_raise_at_once(self) -> None: await asyncio.wait_for(c.connect(ignore), 5) self.assertEqual(self.server.attempts, []) + async def test_outages_with_successes_between_do_not_raise(self) -> None: + c = self.client(subscriptions=["T.AAPL"], max_reconnects=5) + task = self.start(c) + for i in range(1, 8): + await self.server.wait_for_connections(i) + await self.server.wait_for_subscriptions({"T.AAPL"}) + await self.server.drop_connections() + await self.server.wait_for_connections(8) + self.assertFalse(task.done()) + + await c.close() + await asyncio.wait_for(task, 5) + + async def test_consecutive_failures_raise(self) -> None: + c = self.client(max_reconnects=2) + task = self.start(c) + await self.server.wait_for_connections(1) + self.server.reject_handshakes = True + await self.server.drop_connections() + + # the drop and two refused attempts + with self.assertRaises(InvalidStatus): + await asyncio.wait_for(task, 5) + self.assertEqual(len(self.server.attempts), 3) + + async def test_unlimited_reconnects(self) -> None: + c = self.client(max_reconnects=None, max_reconnect_backoff=BACKOFF) + task = self.start(c) + self.server.reject_handshakes = True + await self.server.wait_for(lambda: len(self.server.attempts) >= 20) + self.assertFalse(task.done()) + self.server.reject_handshakes = False + await self.server.wait_for_connections(1) + + await c.close() + await asyncio.wait_for(task, 5) + class TransientTest(unittest.TestCase): """Which failed attempts to connect are retried.""" @@ -267,6 +306,7 @@ def test_defaults(self) -> None: c = WebSocketClient(api_key="key") self.assertEqual(c.reconnect_backoff, 15.0) self.assertEqual(c.max_reconnect_backoff, 120.0) + self.assertIsNone(c.max_reconnects) with mock.patch.object(ws_module.random, "uniform", side_effect=lambda a, b: a): self.assertEqual(self.delays(c, 6), [15, 30, 60, 120, 120, 120]) with mock.patch.object(ws_module.random, "uniform", side_effect=lambda a, b: b): From 4fe66674ce05fea2731dd9f34f919c22c654287a Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:25:01 -0600 Subject: [PATCH 8/9] feat: add a thread-safe request_close() close() is a coroutine that uses the websocket from whatever thread calls it, and while connect() was between connections it only logged "no websocket open to close". request_close() can be called from any thread: it sets the close flag and, through loop.call_soon_threadsafe, wakes connect() and closes the socket, so connect() returns while connected, while connecting, or while waiting to reconnect. close() is kept. A request made while no connect() runs stays pending until a connect() returns because of it, and a stale wake-up queued for an earlier connect() is ignored. If connect() is cancelled while an attempt is in flight, the attempt is cancelled and awaited (websockets bounds it by open_timeout), and a connection it opened meanwhile is closed. --- docs/source/WebSocket.rst | 5 + massive/websocket/__init__.py | 81 ++++++++++++++-- test_websocket/test_reconnect.py | 162 ++++++++++++++++++++++++++++++- 3 files changed, 240 insertions(+), 8 deletions(-) diff --git a/docs/source/WebSocket.rst b/docs/source/WebSocket.rst index c624d9d2..e4a6b412 100644 --- a/docs/source/WebSocket.rst +++ b/docs/source/WebSocket.rst @@ -39,6 +39,11 @@ Close ============================ .. automethod:: massive.WebSocketClient.close +============================ +Request close +============================ +.. automethod:: massive.WebSocketClient.request_close + .. _Stocks getting started: https://massive.com/docs/stocks/ws_getting-started .. _Options getting started: https://massive.com/docs/options/ws_getting-started .. _Forex getting started: https://massive.com/docs/forex/ws_getting-started diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index a651eb99..1fdf4398 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -114,10 +114,11 @@ def __init__( self.json = custom_json else: self.json = json - # set while connect() runs; close() also stops it from reconnecting + # set while connect() runs; close() and request_close() stop it self._loop: Optional[asyncio.AbstractEventLoop] = None self._close_requested = threading.Event() self._close_event: Optional[asyncio.Event] = None + self._close_task: Optional[asyncio.Future] = None self.reconnect_backoff = reconnect_backoff self.max_reconnect_backoff = max_reconnect_backoff @@ -142,7 +143,7 @@ async def connect( than certificate verification), the connection closing during the handshake, and HTTP 408, 429 or 5xx. Other errors, e.g. any other HTTP 4xx or a TLS certificate that cannot be verified, are raised at once. - It returns once `close()` is called. + It returns once `close()` or `request_close()` is called. :param processor: The callback to process messages. :param close_timeout: How long to wait for handshake when calling .close. @@ -158,10 +159,14 @@ async def connect( self._close_event = asyncio.Event() try: await self._connect(processor, close_timeout, **kwargs) + # it returned because close was requested, which is now done + self._close_requested.clear() finally: self._loop = None self._close_event = None - self._close_requested.clear() + closing, self._close_task = self._close_task, None + if closing is not None: + await asyncio.gather(closing, return_exceptions=True) async def _connect( self, @@ -197,7 +202,7 @@ def new_attempt() -> Awaitable[ClientConnection]: if failures > 0 and await self._wait_for_close(delay): return try: - s = await new_attempt() + s = await self._open(new_attempt()) except Exception as e: if not self._is_transient(e): raise @@ -208,8 +213,7 @@ def new_attempt() -> Awaitable[ClientConnection]: delay = self._reconnect_delay(failures - 1) logger.warning("connect failed; reconnecting in %.1f s: %s", delay, e) continue - if self._close_requested.is_set(): - await s.close() + if s is None: return self.websocket = s try: @@ -321,6 +325,40 @@ def _reconnect_delay(self, attempt: int) -> float: ) return random.uniform(base, min(1.2 * base, self.max_reconnect_backoff)) + async def _open( + self, connector: Awaitable[ClientConnection] + ) -> Optional[ClientConnection]: + """Open a connection, or return None if close is requested first.""" + assert self._close_event is not None + attempt = asyncio.ensure_future(connector) + closing = asyncio.ensure_future(self._close_event.wait()) + try: + await asyncio.wait({attempt, closing}, return_when=asyncio.FIRST_COMPLETED) + except BaseException: + # e.g. connect() was cancelled: don't leave a connection open + closing.cancel() + await asyncio.shield(self._abandon(attempt)) + raise + stop = closing.done() or self._close_requested.is_set() + closing.cancel() + if not stop: + return attempt.result() + await self._abandon(attempt) + return None + + @staticmethod + async def _abandon(attempt: "asyncio.Future[ClientConnection]"): + """ + Cancel an attempt to connect, or close the connection it opened. + + Waits for the attempt to end, which websockets bounds by its + `open_timeout`. + """ + attempt.cancel() + await asyncio.wait({attempt}) + if not attempt.cancelled() and attempt.exception() is None: + await attempt.result().close() + async def _wait_for_close(self, timeout: float) -> bool: """Wait up to `timeout` seconds; return whether close was requested.""" assert self._close_event is not None @@ -459,6 +497,37 @@ def unsubscribe_all(self): self.scheduled_subs = set() self.schedule_resub = True + def request_close(self): + """ + Close the websocket connection and make `connect()` return. + + Unlike `close()`, this can be called from any thread, and it does not + wait: `connect()` returns soon after, including while it is + connecting or waiting to reconnect. If no `connect()` is running, the + next one returns without connecting; the request stays pending until a + `connect()` returns because of it. + """ + logger.debug("close requested") + self._close_requested.set() + loop = self._loop + if loop is None: + return + try: + loop.call_soon_threadsafe(self._on_close_requested) + except RuntimeError: + # the loop is closed, so connect() has returned + pass + + def _on_close_requested(self): + # runs on the event loop of connect(); the request may be stale, e.g. + # queued for an earlier connect() on the same loop + if not self._close_requested.is_set(): + return + if self._close_event is not None: + self._close_event.set() + if self.websocket is not None: + self._close_task = asyncio.ensure_future(self.websocket.close()) + async def close(self): """ Close the websocket connection and stop `connect()` from reconnecting. diff --git a/test_websocket/test_reconnect.py b/test_websocket/test_reconnect.py index 33343403..972b2f92 100644 --- a/test_websocket/test_reconnect.py +++ b/test_websocket/test_reconnect.py @@ -1,9 +1,10 @@ import asyncio import socket import ssl +import threading import time import unittest -from typing import Any, List, Optional +from typing import Any, Generator, List, Optional from unittest import mock from fake_server import FakeServer @@ -28,7 +29,7 @@ class TimedServer(FakeServer): """A FakeServer that records when each handshake was attempted.""" - def __init__(self, **kwargs): + def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self.attempts: List[float] = [] # if set, answer handshakes with this HTTP status @@ -47,6 +48,38 @@ async def ignore(msgs: Any) -> None: pass +class FakeConnect: + """ + Stands in for websockets' connect(): each attempt waits for `opened` and + then returns `conn`, even if it was cancelled meanwhile, like a handshake + that completes as it is cancelled. + """ + + def __init__(self) -> None: + self.conn = mock.AsyncMock() + self.attempts = 0 + self.started = asyncio.Event() + self.opened = asyncio.Event() + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + fake = self + + class Attempt: + def __await__(self) -> Generator[Any, None, Any]: + return fake._attempt().__await__() + + return Attempt() + + async def _attempt(self) -> Any: + self.attempts += 1 + self.started.set() + try: + await self.opened.wait() + except asyncio.CancelledError: + pass + return self.conn + + class ReconnectTest(unittest.IsolatedAsyncioTestCase): """How the client behaves when its connection ends.""" @@ -248,6 +281,131 @@ async def test_unlimited_reconnects(self) -> None: await c.close() await asyncio.wait_for(task, 5) + async def request_close_from_a_thread(self, c: WebSocketClient): + thread = threading.Thread(target=c.request_close) + thread.start() + await asyncio.get_running_loop().run_in_executor(None, thread.join) + + async def test_request_close_from_another_thread(self) -> None: + c = self.client(subscriptions=["T.AAPL"]) + task = self.start(c) + await self.server.wait_for_subscriptions({"T.AAPL"}) + # reconnects after a clean close, until request_close() + await self.server.close_connections(code=1000) + await self.server.wait_for_connections(2) + await self.server.wait_for_subscriptions({"T.AAPL"}) + + await self.request_close_from_a_thread(c) + await asyncio.wait_for(task, 2) + await self.server.wait_for(lambda: not self.server.open_connections) + self.assertEqual(len(self.server.connections), 2) + + async def test_request_close_while_waiting_to_reconnect(self) -> None: + c = self.client(reconnect_backoff=60, max_reconnect_backoff=60) + task = self.start(c) + await self.server.wait_for_connections(1) + await self.server.drop_connections() + + await self.request_close_from_a_thread(c) + await asyncio.wait_for(task, 2) + self.assertEqual(len(self.server.attempts), 1) + + async def test_request_close_while_connecting(self) -> None: + # a server that accepts TCP connections but never answers the handshake + accepted = asyncio.Event() + + async def stall(reader: asyncio.StreamReader, writer: asyncio.StreamWriter): + accepted.set() + await reader.read() + writer.close() + + server = await asyncio.start_server(stall, "127.0.0.1", 0) + self.addAsyncCleanup(server.wait_closed) + self.addCleanup(server.close) + port = server.sockets[0].getsockname()[1] + c = WebSocketClient(api_key="key", feed=f"127.0.0.1:{port}", secure=False) + task = self.start(c) + await asyncio.wait_for(accepted.wait(), 5) + + await self.request_close_from_a_thread(c) + await asyncio.wait_for(task, 2) + + async def test_cancel_while_connecting_closes_the_socket(self) -> None: + accepted = asyncio.Event() + disconnected = asyncio.Event() + + async def stall(reader: asyncio.StreamReader, writer: asyncio.StreamWriter): + accepted.set() + await reader.read() + disconnected.set() + writer.close() + + server = await asyncio.start_server(stall, "127.0.0.1", 0) + self.addAsyncCleanup(server.wait_closed) + self.addCleanup(server.close) + port = server.sockets[0].getsockname()[1] + c = WebSocketClient(api_key="key", feed=f"127.0.0.1:{port}", secure=False) + task = asyncio.create_task(c.connect(ignore)) + await asyncio.wait_for(accepted.wait(), 5) + + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + await asyncio.wait_for(disconnected.wait(), 2) + + async def test_cancel_closes_a_connection_opened_meanwhile(self) -> None: + fake = FakeConnect() + with mock.patch.object(ws_module, "connect", fake): + task = asyncio.create_task(self.client().connect(ignore)) + await asyncio.wait_for(fake.started.wait(), 5) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + self.assertEqual(fake.attempts, 1) + fake.conn.close.assert_awaited_once() + + async def test_no_attempt_after_request_close(self) -> None: + fake = FakeConnect() + with mock.patch.object(ws_module, "connect", fake): + c = self.client() + task = asyncio.create_task(c.connect(ignore)) + await asyncio.wait_for(fake.started.wait(), 5) + c.request_close() + # the handshake completes as the close is requested + fake.opened.set() + await asyncio.wait_for(task, 5) + # give a stray attempt every chance to start + for _ in range(20): + await asyncio.sleep(0) + self.assertEqual(fake.attempts, 1) + fake.conn.close.assert_awaited_once() + fake.conn.recv.assert_not_awaited() + + async def test_request_close_before_connect(self) -> None: + c = self.client() + c.request_close() + await asyncio.wait_for(c.connect(ignore), 2) + self.assertEqual(self.server.attempts, []) + + # the request applied to that connect() only + task = self.start(c) + await self.server.wait_for_connections(1) + c.request_close() + await asyncio.wait_for(task, 2) + + async def test_close_request_stays_pending_if_connect_raises(self) -> None: + c = self.client() + c.request_close() + with self.assertRaises(ValueError): + await c.connect(ignore, compression="x") + # the next connect() returns because of it, which consumes it + await asyncio.wait_for(c.connect(ignore), 2) + self.assertEqual(self.server.attempts, []) + task = self.start(c) + await self.server.wait_for_connections(1) + c.request_close() + await asyncio.wait_for(task, 2) + class TransientTest(unittest.TestCase): """Which failed attempts to connect are retried.""" From b2bc14de9d9868850e4c6689083ab715e7e33980 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 16:19:03 -0600 Subject: [PATCH 9/9] feat: surface a max_connections status received during the handshake When the account has too many connections open, the server answers auth with a max_connections status and then closes the connection. The client only logged it, so the caller could not tell it from a network drop. Now connect() sets max_connections_reached until a later handshake succeeds and, in raw mode, passes the message to the processor, as it does for statuses after the handshake. Such a handshake does not reset the failure count, so the backoff keeps growing while the limit is hit. --- massive/websocket/__init__.py | 20 ++++++++- test_websocket/test_reconnect.py | 69 +++++++++++++++++++++++++++++++- 2 files changed, 86 insertions(+), 3 deletions(-) diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 1fdf4398..3dbb2b74 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -119,6 +119,8 @@ def __init__( self._close_requested = threading.Event() self._close_event: Optional[asyncio.Event] = None self._close_task: Optional[asyncio.Future] = None + # whether the server said max_connections during the last handshake + self.max_connections_reached = False self.reconnect_backoff = reconnect_backoff self.max_reconnect_backoff = max_reconnect_backoff @@ -145,6 +147,12 @@ async def connect( 4xx or a TLS certificate that cannot be verified, are raised at once. It returns once `close()` or `request_close()` is called. + If the server sends a `max_connections` status during the handshake + (the account has too many connections open), `max_connections_reached` + is set until a later handshake succeeds, and in raw mode the message is + also passed to `processor`. The server then closes the connection, and + the client reconnects after the backoff. + :param processor: The callback to process messages. :param close_timeout: How long to wait for handshake when calling .close. :param kwargs: Passed to websockets' connect() for each attempt. @@ -219,7 +227,8 @@ def new_attempt() -> Awaitable[ClientConnection]: try: msg = await s.recv() logger.debug("connected: %s", msg) - self._log_server_statuses(self.json.loads(msg)) + msg_parsed = self.json.loads(msg) + self._log_server_statuses(msg_parsed) logger.debug("authing...") await s.send( self.json.dumps({"action": "auth", "params": self.api_key}) @@ -229,7 +238,14 @@ def new_attempt() -> Awaitable[ClientConnection]: logger.debug("authed: %s", auth_msg) if auth_msg_parsed[0]["status"] == "auth_failed": raise AuthError(auth_msg_parsed[0]["message"]) - failures = 0 + self.max_connections_reached = False + for frame, parsed in ((msg, msg_parsed), (auth_msg, auth_msg_parsed)): + if any(m.get("status") == "max_connections" for m in parsed): + self.max_connections_reached = True + if self.raw: + await processor(frame) # type: ignore + if not self.max_connections_reached: + failures = 0 # e.g. max_connections, which the server sends before closing self._log_server_statuses(auth_msg_parsed) self._flush_server_errors(force=True) diff --git a/test_websocket/test_reconnect.py b/test_websocket/test_reconnect.py index 972b2f92..c6cf46d9 100644 --- a/test_websocket/test_reconnect.py +++ b/test_websocket/test_reconnect.py @@ -1,4 +1,5 @@ import asyncio +import json import socket import ssl import threading @@ -7,7 +8,7 @@ from typing import Any, Generator, List, Optional from unittest import mock -from fake_server import FakeServer +from fake_server import FakeServer, status from websockets.asyncio.server import ServerConnection from websockets.datastructures import Headers from websockets.exceptions import ( @@ -454,6 +455,72 @@ def invalid(cause: Optional[BaseException]) -> InvalidMessage: self.assertFalse(WebSocketClient._is_transient(e), repr(e)) +class MaxConnectionsTest(unittest.IsolatedAsyncioTestCase): + """A max_connections status in the reply to auth reaches the caller.""" + + async def asyncSetUp(self) -> None: + self.server = await FakeServer().start() + self.addAsyncCleanup(self.server.stop) + self.status = status( + "max_connections", "Maximum number of websocket connections exceeded." + ) + + async def test_raw_processor_receives_it(self) -> None: + self.server.script_auth(self.status, close_code=1008) + c = WebSocketClient( + api_key="key", + feed=self.server.feed, + secure=False, + raw=True, + reconnect_backoff=BACKOFF, + subscriptions=["T.AAPL"], + ) + frames: List[Any] = [] + + async def processor(msg: Any) -> None: + frames.append((json.loads(msg), c.max_connections_reached)) + + task = asyncio.create_task(c.connect(processor)) + # it reconnects after the server closes, and the next handshake succeeds + await self.server.wait_for_connections(1) + await self.server.wait_for_subscriptions({"T.AAPL"}) + self.assertEqual(len(self.server.connections), 2) + self.assertEqual(frames[0], ([self.status], True)) + self.assertFalse(c.max_connections_reached) + + c.request_close() + await asyncio.wait_for(task, 5) + + async def test_flag_is_set(self) -> None: + self.server.script_auth(self.status, close_code=1008) + self.server.script_auth(self.status, close_code=1008) + c = WebSocketClient( + api_key="key", + feed=self.server.feed, + secure=False, + reconnect_backoff=0.2, + max_reconnect_backoff=0.2, + ) + delays: List[int] = [] + delay = c._reconnect_delay + + def record(attempt: int) -> float: + delays.append(attempt) + return delay(attempt) + + with mock.patch.object(c, "_reconnect_delay", side_effect=record): + task = asyncio.create_task(c.connect(ignore)) + await self.server.wait_for(lambda: c.max_connections_reached) + # cleared by the next successful handshake + await self.server.wait_for_connections(1) + await self.server.wait_for(lambda: not c.max_connections_reached) + self.assertEqual(len(self.server.connections), 3) + c.request_close() + await asyncio.wait_for(task, 5) + # a refused handshake is a failed attempt, so the backoff grows + self.assertEqual(delays, [0, 1]) + + class ReconnectDelayTest(unittest.TestCase): """The wait before each consecutive reconnect attempt."""