From 62e432ff98f74dab756190e7febf06e9cbdedc24 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:20:35 -0600 Subject: [PATCH 01/11] test: stop pook mocks leaking between REST test classes BaseTest turned pook on and registered every mock in setUpClass but never turned it off. Mocks piled up across classes, and pook kept intercepting urllib3 after the REST mock tests had run. That breaks any later test that makes real HTTP requests: pook rejects them with PookNoMatches because it has no mock for them. The next commit adds such a test, which talks to a local HTTP server, and it fails in the full suite without this fix. The leak also hid a problem. Two tests repeat a request, and each mock answers only once, so their second request was served by a copy that an earlier class had registered. Both fail when run on their own. Turn pook off in tearDownClass. Let a test class name the mocks that must answer more than once; those are registered with persist. Other mocks still answer once, because the pagination tests depend on it: a persistent first-page mock would also match the request for the next page. --- test_rest/base.py | 20 +++++++++++++++++++- test_rest/test_dividends.py | 3 +++ test_rest/test_snapshots.py | 3 +++ 3 files changed, 25 insertions(+), 1 deletion(-) diff --git a/test_rest/base.py b/test_rest/base.py index 40854146..bf3d1555 100644 --- a/test_rest/base.py +++ b/test_rest/base.py @@ -3,6 +3,7 @@ import os import re import unittest +from typing import Set import pook # type: ignore # mocks are stored in file tree @@ -32,6 +33,12 @@ class BaseTest(unittest.TestCase): + # Paths (as listed in `mocks`) whose mock answers any number of requests, + # for tests that repeat a request. Every other mock answers once: the + # pagination tests rely on that, because a first-page mock also matches + # the request for the next page. + persistent_mocks: Set[str] = set() + @classmethod def setUpClass(cls): cls.maxDiff = None @@ -40,4 +47,15 @@ def setUpClass(cls): for m in mocks: url = cls.c.BASE + m[0] # print("register", url) - pook.get(url, reply=200, response_body=m[1]) + pook.get( + url, + reply=200, + response_body=m[1], + persist=m[0] in cls.persistent_mocks, + ) + + @classmethod + def tearDownClass(cls): + # stop intercepting urllib3 and drop the mocks, so they don't leak into + # later test classes + pook.off() diff --git a/test_rest/test_dividends.py b/test_rest/test_dividends.py index 7b6575d0..8549df9d 100644 --- a/test_rest/test_dividends.py +++ b/test_rest/test_dividends.py @@ -3,6 +3,9 @@ class DividendsTest(BaseTest): + # test_list_stocks_dividends_ticker_any_of requests this twice + persistent_mocks = {"/stocks/v1/dividends?ticker.any_of=AAPL%2CVZ"} + def test_list_dividends(self): dividends = [d for d in self.c.list_dividends()] expected = [ diff --git a/test_rest/test_snapshots.py b/test_rest/test_snapshots.py index 3ead4e15..4472970a 100644 --- a/test_rest/test_snapshots.py +++ b/test_rest/test_snapshots.py @@ -26,6 +26,9 @@ class SnapshotsTest(BaseTest): + # test_get_snapshot_indices requests this twice + persistent_mocks = {"/v3/snapshot/indices?ticker.any_of=SPX%2CAPx%2CAPy"} + def test_list_universal_snapshots(self): expected = [ UniversalSnapshot( From c08bd5a391182339c43bae00225ebb70dd733a34 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:21:45 -0600 Subject: [PATCH 02/11] test: stress a shared RESTClient from many threads 16 threads share one RESTClient and make 200 calls each against a local HTTP server that echoes every request back. The threads start together behind a barrier, and the server records how many requests it had in flight at once, so the test checks that they really overlapped. Every response is compared with the request that produced it: the path, the exact query string, the Authorization header, and the edge headers. Odd calls send per-request edge headers and even calls send none, so a header leaking from one thread into another call shows up as a mismatch. After the run, the client's default headers must be unchanged: same keys, same values, same order. The calls with headers go through RESTClient._get, which every endpoint method uses. Passing `options` to an endpoint method also adds it to the query string, and that would get in the way of checking the exact query. --- test_rest/test_concurrency.py | 145 ++++++++++++++++++++++++++++++++++ 1 file changed, 145 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..c7def92b --- /dev/null +++ b/test_rest/test_concurrency.py @@ -0,0 +1,145 @@ +import json +import threading +import unittest +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any, Dict, List, Optional, Tuple +from urllib.parse import parse_qsl, urlparse + +from massive import RESTClient +from massive.rest.models.request import ( + RequestOptionBuilder, + X_MASSIVE_EDGE_ID, + X_MASSIVE_EDGE_IP_ADDRESS, +) + + +class EchoServer(ThreadingHTTPServer): + """Answers every GET with what it received. + + It counts the requests and records the most requests it had in flight at + once, which shows whether the client's threads really overlapped. + """ + + daemon_threads = True + + def __init__(self): + super().__init__(("127.0.0.1", 0), EchoHandler) + self.lock = threading.Lock() + self.count = 0 + self.in_flight = 0 + self.peak_in_flight = 0 + + +class EchoHandler(BaseHTTPRequestHandler): + # keep-alive, so the client reuses pooled connections across threads + protocol_version = "HTTP/1.1" + server: EchoServer + + def do_GET(self): + with self.server.lock: + self.server.count += 1 + self.server.in_flight += 1 + self.server.peak_in_flight = max( + self.server.peak_in_flight, self.server.in_flight + ) + try: + url = urlparse(self.path) + body = json.dumps( + { + "path": url.path, + "query": dict(parse_qsl(url.query)), + "edge_id": self.headers.get(X_MASSIVE_EDGE_ID), + "edge_ip": self.headers.get(X_MASSIVE_EDGE_IP_ADDRESS), + "authorization": self.headers.get("Authorization"), + } + ).encode() + 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) + finally: + with self.server.lock: + self.server.in_flight -= 1 + + 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 + + def setUp(self): + self.httpd = EchoServer() + thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) + thread.start() + self.addCleanup(thread.join) + self.addCleanup(self.httpd.server_close) + self.addCleanup(self.httpd.shutdown) + host, port = self.httpd.server_address[:2] + self.base = f"http://{host!s}:{port}" + + def test_shared_client(self): + c = RESTClient("key", base=self.base, retries=0, maxsize=self.threads) + headers_before = list(c.headers.items()) + barrier = threading.Barrier(self.threads, timeout=10) + + def run(t: int) -> List[Tuple[Dict[str, Any], Dict[str, Any]]]: + # start every thread together, so their requests overlap + barrier.wait() + results = [] + for i in range(self.calls): + ticker = f"T{t}N{i}" + date = f"2024-01-{t + 1:02}" + edge_id: Optional[str] = None + edge_ip: Optional[str] = None + if i % 2: + # Odd calls send per-request headers. They go through _get, + # the method every endpoint calls, because the endpoint + # methods also copy `options` into the query string. + edge_id, edge_ip = f"edge-{t}-{i}", f"10.0.{t}.{i % 256}" + resp = c._get( + path=f"/v3/reference/tickers/{ticker}", + params={"ticker": ticker, "date": date}, + raw=True, + options=RequestOptionBuilder( + edge_id=edge_id, edge_ip_address=edge_ip + ), + ) + else: + # raw=True returns the response, so the echoed request can + # be checked instead of a deserialized TickerDetails + resp = c.get_ticker_details(ticker, date=date, raw=True) + expected = { + "path": f"/v3/reference/tickers/{ticker}", + "query": {"ticker": ticker, "date": date}, + "edge_id": edge_id, + "edge_ip": edge_ip, + "authorization": "Bearer key", + } + results.append((expected, json.loads(resp.data))) + 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()] + + # every response answers its own request: the path and query name the + # thread and call, and the edge headers are the ones that call sent + 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) + # the threads really shared the client at the same time + self.assertGreater(self.httpd.peak_in_flight, 1) + # per-request headers must not leak into the client's default headers: + # same keys, same values, same order + self.assertEqual(list(c.headers.items()), headers_before) + + +if __name__ == "__main__": + unittest.main() From 4e5d32d4f29eb38f63930a9fd5e34ec565d16755 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 03/11] 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. --- 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 5fbf82db5d71c029a792ac51bf74b4efb43fabdd 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 04/11] 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. --- .github/workflows/test.yml | 40 +++++++++++++++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index a1ae9ada..48ef68c9 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,41 @@ 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 - <<'EOF' + import sys, sysconfig + state = ( + f"version={sys.version!r} " + f"Py_GIL_DISABLED={sysconfig.get_config_var('Py_GIL_DISABLED')!r} " + f"_is_gil_enabled={sys._is_gil_enabled()!r}" + ) + print(state) + assert not sys._is_gil_enabled(), f"the GIL is enabled: {state}" + EOF + make test From 3adbdc312cc3c1321f4452d7dcf0e3f7cc2657b5 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:32:36 -0600 Subject: [PATCH 05/11] test: add a scriptable fake Massive WebSocket server The WebSocket tests use mock_server, which listens on a fixed port, keeps module-level state and cannot script failures, so they cannot check what the client sends or how it behaves when the server misbehaves. FakeServer listens on a free port and records, per connection, every message the client sent and the topics it is subscribed to. It speaks the Massive protocol by default and can script the reply to auth (for example a max_connections status followed by a close), refuse handshakes, close or drop connections, and run a hook on each client message, which may itself close or drop the connection. Its waits are driven by changes to its state, and its state is only touched from its own event loop. --- test_websocket/fake_server.py | 380 +++++++++++++++++++++++++++++ test_websocket/test_fake_server.py | 39 +++ 2 files changed, 419 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..6561319d --- /dev/null +++ b/test_websocket/fake_server.py @@ -0,0 +1,380 @@ +""" +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 assert on exactly what a client sent. Run it inside the +test's event loop:: + + server = await FakeServer().start() + 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 + ... + await server.stop() + +By default the server speaks the Massive protocol: it sends ``connected``, +answers ``auth`` with ``auth_success``, and acknowledges each subscribed or +unsubscribed topic with a ``success`` status. Scripted replies and hooks +override that behaviour. +""" + +import asyncio +import json +import threading +from dataclasses import dataclass +from http import HTTPStatus +from typing import Any, Awaitable, Callable, Dict, List, Optional, Set + +from websockets.asyncio.server import Server, ServerConnection, serve +from websockets.exceptions import ConnectionClosed +from websockets.http11 import Request, Response + + +def status(kind: str, message: str) -> Dict[str, Any]: + """Build a Massive status message, e.g. ``status("max_connections", "...")``.""" + return {"ev": "status", "status": kind, "message": message} + + +@dataclass +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. + """ + + frames: List[Dict[str, Any]] + close_code: Optional[int] = None + close_reason: str = "" + + +class FakeConnection: + """ + One client connection accepted by a ``FakeServer``. + + Like the server, it belongs to the server's event loop: read its attributes + and call its methods only from that loop. + """ + + def __init__(self, websocket: ServerConnection): + self.websocket = websocket + 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() + # the task running the server's handler for this connection + self.handler: Optional["asyncio.Task[Any]"] = None + + @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 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._wait_closed() + + 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._wait_closed() + + async def _wait_closed(self): + # The handler sets closed when it returns. Called from the handler + # itself, e.g. from an on_message hook, waiting would never end: the + # handler returns once the hook does, as the connection is closed. + if asyncio.current_task() is not self.handler: + await self.closed.wait() + + +class FakeServer: + """ + A fake Massive WebSocket server listening on ``host`` and a free port. + + Threading model: the server and all its state (``connections``, each + connection's ``inbox`` and ``subscriptions``, the scripted replies) live on + the event loop that called ``start()``, and are only read or changed there. + Clients under test may run on other threads with their own event loops; they + talk to the server only over the socket. The state accessors raise + ``RuntimeError`` when called from another thread, and return copies, so a + test never holds a collection the server is changing. + + :param host: The interface to listen on. + :param ack_subscriptions: Whether to acknowledge subscribe/unsubscribe. + :param on_message: Called on the server's loop with every decoded client + message after it has been recorded and handled; use it to script replies. + It may close or drop the connection. + """ + + def __init__( + self, + host: str = "127.0.0.1", + ack_subscriptions: bool = True, + on_message: Optional[Callable[[FakeConnection, Any], Awaitable[None]]] = None, + ): + self.host = host + self.port = 0 + self.ack_subscriptions = ack_subscriptions + 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 + self._thread: Optional[int] = None + self._changed: Optional[asyncio.Condition] = None + + # lifecycle + + async def start(self) -> "FakeServer": + self._thread = threading.get_ident() + self._changed = asyncio.Condition() + 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 + + @property + def feed(self) -> str: + """The ``feed`` to pass to ``WebSocketClient`` (with ``secure=False``).""" + return f"{self.host}:{self.port}" + + def _check_thread(self): + if self._thread is not None and threading.get_ident() != self._thread: + raise RuntimeError( + "FakeServer state must be used from the thread running its event loop" + ) + + # 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)``. A + reply without ``auth_success`` leaves the connection unauthenticated, + and one without frames sends nothing. + """ + self._check_thread() + self._auth_replies.append(AuthReply(list(frames), close_code, close_reason)) + + @property + def connections(self) -> List[FakeConnection]: + """Every connection accepted so far, oldest first.""" + self._check_thread() + return list(self._connections) + + @property + def connection(self) -> FakeConnection: + """The most recent connection.""" + self._check_thread() + if not self._connections: + raise AssertionError("no client has connected") + return self._connections[-1] + + @property + def open_connections(self) -> List[FakeConnection]: + self._check_thread() + 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.""" + self._check_thread() + 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.""" + self._check_thread() + return set(self._connections[-1].subscriptions) if self._connections else set() + + async def send_data(self, events: List[Dict[str, Any]]): + """Send data events, e.g. ``[{"ev": "T", "sym": "AAPL", ...}]``, to every + open, authenticated connection.""" + for c in self.open_connections: + if c.authenticated: + await c.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, + description: str = "condition", + ): + """ + Wait until ``predicate()`` is true. + + The predicate is checked now and after every change to the server's + state, so it must depend only on that state. On timeout this fails with + an ``AssertionError`` naming ``description`` and the server's state. + """ + self._check_thread() + if self._changed is None: + raise RuntimeError("FakeServer is not started") + changed = self._changed + + async def wait(): + async with changed: + await changed.wait_for(predicate) + + try: + await asyncio.wait_for(wait(), timeout) + except asyncio.TimeoutError: + raise AssertionError( + f"timed out after {timeout}s waiting for {description}; {self._state()}" + ) from None + + 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, + f"{n} authenticated connections", + ) + + 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, + f"subscriptions {sorted(topics)}", + ) + + 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, "a matching client message") + return found[0] + + def _state(self) -> str: + return ( + f"connections={len(self._connections)} " + f"open={len(self.open_connections)} " + f"subscriptions={sorted(self.subscriptions)} " + f"last_messages={self.inbox[-5:]}" + ) + + async def _notify(self): + if self._changed is not None: + async with self._changed: + self._changed.notify_all() + + # 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(websocket) + conn.handler = asyncio.current_task() + self._connections.append(conn) + await self._notify() + 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) + await self._notify() + except ConnectionClosed: + pass + finally: + conn.closed.set() + await self._notify() + + 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 + 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": + 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..f1b9251c --- /dev/null +++ b/test_websocket/test_fake_server.py @@ -0,0 +1,39 @@ +import asyncio +import unittest +from typing import Any + +from fake_server import FakeConnection, FakeServer +from websockets.asyncio.client import connect +from websockets.exceptions import ConnectionClosed + + +class FakeServerTest(unittest.IsolatedAsyncioTestCase): + """The fake server's own behaviour, where tests rely on it.""" + + async def hook_closes(self, close: str) -> None: + async def on_message(conn: FakeConnection, msg: Any): + if isinstance(msg, dict) and msg.get("action") == "subscribe": + if close == "drop": + await conn.drop() + else: + await conn.close(1011) + + server = await FakeServer(on_message=on_message).start() + async with connect(f"ws://{server.feed}/stocks") as ws: + await ws.recv() # connected + await ws.send('{"action": "subscribe", "params": "T.A"}') + with self.assertRaises(ConnectionClosed): + while True: + await asyncio.wait_for(ws.recv(), 5) + await server.wait_for(lambda: not server.connection.open, 5) + await asyncio.wait_for(server.stop(), 5) + + async def test_hook_can_close_the_connection(self) -> None: + await self.hook_closes("close") + + async def test_hook_can_drop_the_connection(self) -> None: + await self.hook_closes("drop") + + +if __name__ == "__main__": + unittest.main() From f0949c2197b2f12590f609a9381407fcda7dac84 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:22:20 -0600 Subject: [PATCH 06/11] fix: make WebSocket subscribe and unsubscribe thread-safe subscribe(), unsubscribe() and unsubscribe_all() changed scheduled_subs and schedule_resub on the caller's thread while connect() read them on the event loop's thread, with no synchronization. A change made between the loop computing the frames and clearing schedule_resub was lost, so the server could end up with a different set than the one requested. A lock now guards both attributes, both where the callers change them and where connect() resets them after an abnormal close; the loop copies the scheduled set and clears the flag under it, and sends the frames from the copy. The methods replace scheduled_subs with a new set instead of changing it in place, so another thread can read or iterate it without the lock. It remains a set, and the frames are sent in the same order as before. --- massive/websocket/__init__.py | 162 ++++++++++++++++++--------- test_websocket/test_subscriptions.py | 83 ++++++++++++++ 2 files changed, 194 insertions(+), 51 deletions(-) create mode 100644 test_websocket/test_subscriptions.py diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 7463b4cb..7ad934aa 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -1,11 +1,22 @@ import os from collections import Counter from enum import Enum -from typing import Optional, Union, List, Set, Callable, Awaitable, Any +from typing import ( + Optional, + Union, + List, + Set, + Callable, + Awaitable, + Any, + Iterable, + Tuple, +) import logging import json import asyncio import ssl +import threading import time from .models import * from websockets.asyncio.client import connect, ClientConnection @@ -76,8 +87,14 @@ def __init__( self.websocket: Optional[ClientConnection] = None if subscriptions is None: subscriptions = [] + # replaced on every change rather than changed in place, so other + # threads can read it self.scheduled_subs: Set[str] = set(subscriptions) self.schedule_resub = True + # guards scheduled_subs and schedule_resub, which subscribe() and + # friends change on the caller's thread while connect() reads them on + # the event loop's thread + self._subs_lock = threading.Lock() self._server_errors: Counter = Counter() self._server_errors_since = 0.0 self.ca_certs = ca_certs @@ -133,17 +150,11 @@ async def connect( self._log_server_statuses(auth_msg_parsed) self._flush_server_errors(force=True) while True: + # read without the lock to keep the per-message cost low; + # _reconcile() checks again under it, and a change missed + # here is seen on the next iteration if self.schedule_resub: - logger.debug( - "reconciling: %s %s", self.subs, self.scheduled_subs - ) - new_subs = self.scheduled_subs.difference(self.subs) - await self._subscribe(new_subs) - old_subs = self.subs.difference(self.scheduled_subs) - await self._unsubscribe(old_subs) - self.subs = self.scheduled_subs - self.subs = set(self.scheduled_subs) - self.schedule_resub = False + await self._reconcile(s) try: cmsg: Union[List[WebSocketMessage], Union[str, bytes]] = ( @@ -175,9 +186,10 @@ async def connect( logger.debug("connection closed (ERR): %s", e) self._flush_server_errors(force=True) reconnects += 1 - self.scheduled_subs = set(self.subs) - self.subs = set() - self.schedule_resub = True + with self._subs_lock: + self.scheduled_subs = set(self.subs) + self.subs = set() + self.schedule_resub = True if self.max_reconnects is not None and reconnects > self.max_reconnects: break continue @@ -242,23 +254,49 @@ def _flush_server_errors(self, force: bool = False): ) self._server_errors = Counter() - async def _subscribe(self, topics: Union[List[str], Set[str]]): - if self.websocket is None or len(topics) == 0: + async def _subscribe( + self, + topics: Union[List[str], Set[str]], + ws: Optional[ClientConnection] = None, + ): + if ws is None: + ws = self.websocket + if ws is None or len(topics) == 0: return subs = ",".join(topics) logger.debug("subbing: %s", subs) - await self.websocket.send( - self.json.dumps({"action": "subscribe", "params": subs}) - ) + await ws.send(self.json.dumps({"action": "subscribe", "params": subs})) - async def _unsubscribe(self, topics: Union[List[str], Set[str]]): - if self.websocket is None or len(topics) == 0: + async def _unsubscribe( + self, + topics: Union[List[str], Set[str]], + ws: Optional[ClientConnection] = None, + ): + if ws is None: + ws = self.websocket + if ws is None or len(topics) == 0: return subs = ",".join(topics) logger.debug("unsubbing: %s", subs) - await self.websocket.send( - self.json.dumps({"action": "unsubscribe", "params": subs}) - ) + await ws.send(self.json.dumps({"action": "unsubscribe", "params": subs})) + + async def _reconcile(self, ws: ClientConnection): + """Send ``ws`` the frames that move it from ``subs`` to ``scheduled_subs``.""" + with self._subs_lock: + if not self.schedule_resub: + return + desired = set(self.scheduled_subs) + self.schedule_resub = False + logger.debug("reconciling: %s %s", self.subs, desired) + try: + await self._subscribe(desired.difference(self.subs), ws) + await self._unsubscribe(self.subs.difference(desired), ws) + except BaseException: + # e.g. the connection closed: try again on the next connection + with self._subs_lock: + self.schedule_resub = True + raise + self.subs = desired @staticmethod def _parse_subscription(s: str): @@ -270,53 +308,75 @@ def _parse_subscription(s: str): return split + def _parse_subscriptions( + self, subscriptions: Iterable[str] + ) -> List[Tuple[str, str, str]]: + """``(topic, channel, symbol)`` per well-formed topic; logs the rest.""" + parsed = [] + for s in subscriptions: + topic, sym = self._parse_subscription(s) + if topic is not None: + parsed.append((s, topic, sym)) + return parsed + def subscribe(self, *subscriptions: str): """ Subscribe to given subscriptions. + Safe to call from any thread, before or while ``connect()`` runs. + :param subscriptions: Subscriptions (args) """ - for s in subscriptions: - topic, sym = self._parse_subscription(s) - if topic == None: - continue + parsed = self._parse_subscriptions(subscriptions) + with self._subs_lock: + scheduled = set(self.scheduled_subs) + for s, topic, sym in parsed: + scheduled.add(s) + # If user subs to X.*, remove other X.\w+ + if sym == "*": + for t in list(self.subs): + if t.startswith(topic): + scheduled.discard(t) + + self.scheduled_subs = scheduled + self.schedule_resub = True + for s, _, _ in parsed: logger.debug("sub desired: %s", s) - self.scheduled_subs.add(s) - # If user subs to X.*, remove other X.\w+ - if sym == "*": - for t in list(self.subs): - if t.startswith(topic): - self.scheduled_subs.discard(t) - - self.schedule_resub = True def unsubscribe(self, *subscriptions: str): """ Unsubscribe from given subscriptions. + Safe to call from any thread, before or while ``connect()`` runs. + :param subscriptions: Subscriptions (args) """ - for s in subscriptions: - topic, sym = self._parse_subscription(s) - if topic == None: - continue + parsed = self._parse_subscriptions(subscriptions) + with self._subs_lock: + scheduled = set(self.scheduled_subs) + for s, topic, sym in parsed: + scheduled.discard(s) + + # If user unsubs to X.*, remove other X.\w+ + if sym == "*": + for t in list(self.subs): + if t.startswith(topic): + scheduled.discard(t) + + self.scheduled_subs = scheduled + self.schedule_resub = True + for s, _, _ in parsed: logger.debug("sub undesired: %s", s) - self.scheduled_subs.discard(s) - - # If user unsubs to X.*, remove other X.\w+ - if sym == "*": - for t in list(self.subs): - if t.startswith(topic): - self.scheduled_subs.discard(t) - - self.schedule_resub = True def unsubscribe_all(self): """ Unsubscribe from all subscriptions. + + Safe to call from any thread, before or while ``connect()`` runs. """ - self.scheduled_subs = set() - self.schedule_resub = True + with self._subs_lock: + self.scheduled_subs = set() + self.schedule_resub = True async def close(self): """ diff --git a/test_websocket/test_subscriptions.py b/test_websocket/test_subscriptions.py new file mode 100644 index 00000000..72c1050b --- /dev/null +++ b/test_websocket/test_subscriptions.py @@ -0,0 +1,83 @@ +import asyncio +import threading +import unittest +from concurrent.futures import ThreadPoolExecutor +from typing import Dict, Set + +from fake_server import FakeServer + +from massive import WebSocketClient + + +async def ignore(msgs): + pass + + +class SubscriptionsTest(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 start(self, c: WebSocketClient) -> "asyncio.Task[None]": + """Run ``c.connect()`` until the test ends; wait for it to connect.""" + task = asyncio.create_task(c.connect(ignore)) + + async def stop(): + if not task.done(): + await c.close() + await asyncio.wait_for(task, 5) + + self.addAsyncCleanup(stop) + await self.server.wait_for_connections(len(self.server.connections) + 1) + return task + + +class ThreadSafetyTest(SubscriptionsTest): + threads = 16 + calls = 10_000 + topics_per_thread = 50 + + async def test_concurrent_subscribe_unsubscribe(self) -> None: + # without acknowledgements, so the client does not have tens of + # thousands of them to read before it can close + self.server.ack_subscriptions = False + c = self.client(subscriptions=["T.INITIAL"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.INITIAL"}) + barrier = threading.Barrier(self.threads, timeout=10) + # the last call for each topic decides whether it is subscribed + last: Dict[str, bool] = {} + last_lock = threading.Lock() + + def run(i: int): + mine: Dict[str, bool] = {} + barrier.wait() + for j in range(self.calls): + topic = f"T.S{i}X{j % self.topics_per_thread}" + if j % 3 == 0: + c.unsubscribe(topic) + mine[topic] = False + else: + c.subscribe(topic) + mine[topic] = True + with last_lock: + last.update(mine) + + 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)) + ) + + desired: Set[str] = {"T.INITIAL"} | {t for t, on in last.items() if on} + self.assertEqual(c.scheduled_subs, desired) + await self.server.wait_for_subscriptions(desired, timeout=10) + self.assertEqual(len(self.server.open_connections), 1) + + +if __name__ == "__main__": + unittest.main() From c23d9fb9cd6044b5c4f4135ce649ac27b6e38025 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:27:14 -0600 Subject: [PATCH 07/11] fix: send WebSocket subscription changes without waiting for recv connect() applied scheduled subscription changes only when its receive loop came round again, after a message or the one-second recv timeout, so on a quiet feed a subscribe() took up to a second to reach the server. Each connection now gets one sender task, started after authentication. It waits on an asyncio.Event and, when woken, sends the frames that move the server from subs to the latest scheduled set; it starts by sending the whole set, so a reconnect re-sends the desired state. Being the only sender, it never interleaves frames, and a burst of changes is merged into one round. subscribe(), unsubscribe() and unsubscribe_all() change the scheduled set under the lock and then, if the change is the first since the sender last read it, set the event through call_soon_threadsafe, which never blocks the caller. The sender sends on the socket it was started for, not self.websocket. When the connection ends, the sender is cancelled and awaited before subs is reset. An error other than ConnectionClosed ends the sender and is raised by connect() at the next turn of the receive loop; a closed connection is left to recv() to report, after the messages already received are read. The receive loop still looks at schedule_resub at the top of each iteration, now to wake the sender, so code that changes scheduled_subs itself and sets the flag keeps working. Otherwise the receive loop is unchanged. Sending while another coroutine waits in recv() on the same connection is safe in websockets 14.0 to 15.0.1 (websockets/asyncio/connection.py): recv() only reads from the message assembler (14.0 L270, 14.2 L294, 15.0.1 L303), whose only concurrency check rejects a second concurrent reader (asyncio/messages.py, Assembler.get: 14.0 L142, 14.2 L141, 15.0.1 L150), and send() waits only for a fragmented send in progress (14.0 L427, 14.2 L451, 15.0.1 L469) before writing through send_context() (14.0 L434, 14.2 L458, 15.0.1 L476). --- massive/websocket/__init__.py | 132 ++++++++++++---- test_websocket/test_subscriptions.py | 222 ++++++++++++++++++++++++++- 2 files changed, 323 insertions(+), 31 deletions(-) diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 7ad934aa..18bcf491 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -20,7 +20,11 @@ import time from .models import * from websockets.asyncio.client import connect, ClientConnection -from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError +from websockets.exceptions import ( + ConnectionClosed, + ConnectionClosedOK, + ConnectionClosedError, +) from ..certs import ca_bundle_path from ..logging import get_logger import logging @@ -95,6 +99,9 @@ def __init__( # friends change on the caller's thread while connect() reads them on # the event loop's thread self._subs_lock = threading.Lock() + # while a connection is open: its loop, and the event that wakes its + # sender task (see _start_sender) + self._wake: Optional[Tuple[asyncio.AbstractEventLoop, asyncio.Event]] = None self._server_errors: Counter = Counter() self._server_errors_since = 0.0 self.ca_certs = ca_certs @@ -149,33 +156,40 @@ async def connect( # e.g. max_connections, which the server sends before closing self._log_server_statuses(auth_msg_parsed) self._flush_server_errors(force=True) - while True: - # read without the lock to keep the per-message cost low; - # _reconcile() checks again under it, and a change missed - # here is seen on the next iteration - if self.schedule_resub: - await self._reconcile(s) - - try: - cmsg: Union[List[WebSocketMessage], Union[str, bytes]] = ( - await asyncio.wait_for(s.recv(), timeout=1) - ) - except asyncio.TimeoutError: + # subscription frames go out from a task of their own, starting + # with the whole scheduled set + sender, wake = self._start_sender(s) + try: + while True: + if sender.done(): + self._check_sender(sender) + if self.schedule_resub: + # e.g. set by code that changed scheduled_subs + # itself, as the loop used to check it here + wake.set() + + try: + cmsg: Union[List[WebSocketMessage], Union[str, bytes]] = ( + await asyncio.wait_for(s.recv(), timeout=1) + ) + except asyncio.TimeoutError: + self._flush_server_errors() + continue + + if not self.raw: + # we know cmsg is Data + msgJson = self.json.loads(cmsg) # type: ignore + for m in msgJson: + if m["ev"] == "status": + logger.debug("status: %s", m["message"]) + self._log_server_statuses(msgJson) + cmsg = parse(msgJson, logger, self.market) self._flush_server_errors() - continue - - if not self.raw: - # we know cmsg is Data - msgJson = self.json.loads(cmsg) # type: ignore - for m in msgJson: - if m["ev"] == "status": - logger.debug("status: %s", m["message"]) - self._log_server_statuses(msgJson) - cmsg = parse(msgJson, logger, self.market) - self._flush_server_errors() - - if len(cmsg) > 0: - await processor(cmsg) # type: ignore + + if len(cmsg) > 0: + await processor(cmsg) # type: ignore + finally: + await self._stop_sender(sender) except ConnectionClosedOK as e: logger.debug("connection closed (OK): %s", e) # the server sends a status (e.g. slow_consumer) before closing @@ -280,6 +294,58 @@ async def _unsubscribe( logger.debug("unsubbing: %s", subs) await ws.send(self.json.dumps({"action": "unsubscribe", "params": subs})) + def _start_sender( + self, ws: ClientConnection + ) -> Tuple["asyncio.Task[None]", asyncio.Event]: + """Start the task that sends ``ws`` subscription changes; see _notify.""" + wake = asyncio.Event() + wake.set() # first, the whole scheduled set + self._wake = (asyncio.get_running_loop(), wake) + return asyncio.ensure_future(self._send_changes(ws, wake)), wake + + async def _send_changes(self, ws: ClientConnection, wake: asyncio.Event): + # The only sender of subscription frames on ws, so frames for two + # changes never interleave. Sending while connect() waits in recv() is + # safe: in websockets, only two concurrent recv() calls conflict. + while True: + await wake.wait() + wake.clear() + await self._reconcile(ws) + + def _check_sender(self, sender: "asyncio.Task[None]"): + # raise what ended the sender, unless the connection closed: then + # recv() raises too, once the messages already received are read + if sender.cancelled(): + return + exc = sender.exception() + if exc is not None and not isinstance(exc, ConnectionClosed): + raise exc + + async def _stop_sender(self, sender: "asyncio.Task[None]"): + self._wake = None + sender.cancel() + await asyncio.wait({sender}) # does not raise + self._check_sender(sender) + + def _request_resub(self) -> bool: + # Call with the lock held. Only a change that sets the flag needs to + # wake the sender: while it is set, the sender has been woken already, + # or there is none and the next one starts by reconciling. + wake = not self.schedule_resub + self.schedule_resub = True + return wake + + def _notify(self): + """Wake the sender, if a connection is open. Safe from any thread.""" + waker = self._wake + if waker is None: + return + loop, wake = waker + try: + loop.call_soon_threadsafe(wake.set) + except RuntimeError: + pass # the loop is closed + async def _reconcile(self, ws: ClientConnection): """Send ``ws`` the frames that move it from ``subs`` to ``scheduled_subs``.""" with self._subs_lock: @@ -339,7 +405,9 @@ def subscribe(self, *subscriptions: str): scheduled.discard(t) self.scheduled_subs = scheduled - self.schedule_resub = True + wake = self._request_resub() + if wake: + self._notify() for s, _, _ in parsed: logger.debug("sub desired: %s", s) @@ -364,7 +432,9 @@ def unsubscribe(self, *subscriptions: str): scheduled.discard(t) self.scheduled_subs = scheduled - self.schedule_resub = True + wake = self._request_resub() + if wake: + self._notify() for s, _, _ in parsed: logger.debug("sub undesired: %s", s) @@ -376,7 +446,9 @@ def unsubscribe_all(self): """ with self._subs_lock: self.scheduled_subs = set() - self.schedule_resub = True + wake = self._request_resub() + if wake: + self._notify() async def close(self): """ diff --git a/test_websocket/test_subscriptions.py b/test_websocket/test_subscriptions.py index 72c1050b..4d4d3012 100644 --- a/test_websocket/test_subscriptions.py +++ b/test_websocket/test_subscriptions.py @@ -1,12 +1,15 @@ import asyncio +import json import threading import unittest from concurrent.futures import ThreadPoolExecutor -from typing import Dict, Set +from typing import Any, Dict, List, Set +from unittest import mock from fake_server import FakeServer from massive import WebSocketClient +from massive import websocket as websocket_module async def ignore(msgs): @@ -79,5 +82,222 @@ def run(i: int): self.assertEqual(len(self.server.open_connections), 1) +class CompatibilityTest(SubscriptionsTest): + async def test_changing_the_attributes_directly(self) -> None: + # as before, changing scheduled_subs and setting schedule_resub works + c = self.client(subscriptions=["T.A"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.A"}) + c.scheduled_subs.add("T.B") + c.schedule_resub = True + await self.server.wait_for_subscriptions({"T.A", "T.B"}) + self.assertIsInstance(c.scheduled_subs, set) + self.assertIsInstance(c.subs, set) + + async def test_subscribe_after_changing_the_attributes_directly(self) -> None: + c = self.client(subscriptions=["T.A"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.A"}) + c.scheduled_subs.add("T.B") + c.schedule_resub = True + c.subscribe("T.C") # schedule_resub is already set + await self.server.wait_for_subscriptions({"T.A", "T.B", "T.C"}) + c.subscribe("T.D") + await self.server.wait_for_subscriptions({"T.A", "T.B", "T.C", "T.D"}) + + +def recv_timeout(seconds: float): + """ + Patch the client's receive loop to wait up to ``seconds`` for a message, + instead of one second, by giving its module an ``asyncio`` whose + ``wait_for()`` uses that timeout. + """ + + class Asyncio: + def __getattr__(self, name: str) -> Any: + return getattr(asyncio, name) + + @staticmethod + def wait_for(aw: Any, timeout: float) -> Any: + return asyncio.wait_for(aw, seconds) + + return mock.patch.object(websocket_module, "asyncio", Asyncio()) + + +class WakeUpTest(SubscriptionsTest): + """ + With a 30 second receive timeout, the receive loop sees a change only + after the next message: changes reach the server promptly only if they + are sent without waiting for it. + """ + + latency = 0.5 + + async def asyncSetUp(self) -> None: + await super().asyncSetUp() + patcher = recv_timeout(30.0) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_subscribe_is_sent_promptly(self) -> None: + c = self.client(subscriptions=["T.AAPL"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.AAPL"}) + c.subscribe("T.MSFT") + await self.server.wait_for_subscriptions({"T.AAPL", "T.MSFT"}, self.latency) + + async def test_unsubscribe_from_another_thread_is_sent_promptly(self) -> None: + c = self.client(subscriptions=["T.AAPL", "T.MSFT"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.AAPL", "T.MSFT"}) + await asyncio.to_thread(c.unsubscribe, "T.AAPL") + await self.server.wait_for_subscriptions({"T.MSFT"}, self.latency) + + async def test_unsubscribe_all_is_sent_promptly(self) -> None: + c = self.client(subscriptions=["T.AAPL"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.AAPL"}) + await asyncio.to_thread(c.unsubscribe_all) + await self.server.wait_for_subscriptions(set(), self.latency) + + async def test_changes_between_connect_calls(self) -> None: + c = self.client() + c.subscribe("T.AAPL") # before connect() runs + + async def processor(msgs): + await c.close() + + # connect() in a thread with its own event loop, closed afterwards + connected = asyncio.create_task( + asyncio.to_thread(asyncio.run, c.connect(processor)) + ) + await self.server.wait_for_subscriptions({"T.AAPL"}) + await self.server.send_data( + [{"ev": "T", "sym": "AAPL", "x": 1, "p": 1, "s": 1, "t": 1}] + ) + await asyncio.wait_for(connected, 5) + + # connect()'s loop is closed: scheduling still works and does not raise + c.subscribe("T.MSFT") + c.unsubscribe("T.AAPL") + self.assertEqual(c.scheduled_subs, {"T.MSFT"}) + + async def test_frequent_receive_timeouts(self) -> None: + errors: List[Dict[str, Any]] = [] + loop = asyncio.get_running_loop() + handler = loop.get_exception_handler() + loop.set_exception_handler(lambda loop, context: errors.append(context)) + self.addCleanup(loop.set_exception_handler, handler) + + c = self.client(subscriptions=["T.AAPL"]) + with recv_timeout(0.0005): + await self.start(c) + await self.server.wait_for_subscriptions({"T.AAPL"}) + + def hammer(): + for _ in range(20_000): + c.subscribe("T.X") + c.unsubscribe("T.X") + + await asyncio.to_thread(hammer) + await self.server.wait_for_subscriptions({"T.AAPL"}) + self.assertEqual(errors, []) + + async def test_changes_while_receiving_data(self) -> None: + # subscription frames go out while the receive loop is in recv() and + # the server is streaming data to it + errors: List[Dict[str, Any]] = [] + loop = asyncio.get_running_loop() + handler = loop.get_exception_handler() + loop.set_exception_handler(lambda loop, context: errors.append(context)) + self.addCleanup(loop.set_exception_handler, handler) + + n = 300 + received: List[str] = [] + all_received = asyncio.Event() + + async def processor(msgs): + received.extend(m.symbol for m in msgs) + if len(received) >= n: + all_received.set() + + c = self.client(subscriptions=["T.AAPL"]) + task = asyncio.create_task(c.connect(processor)) + await self.server.wait_for_subscriptions({"T.AAPL"}) + subscribed = {"T.AAPL"} + for i in range(n): + await self.server.send_data( + [{"ev": "T", "sym": f"S{i}", "x": 1, "p": 1, "s": 1, "t": 1}] + ) + if i % 30 == 29: + # a burst of data is on its way to the client: change the + # subscriptions, alternately from this thread and another + topic = f"Q.S{i}" + if i % 60 == 29: + c.subscribe(topic) + else: + await asyncio.to_thread(c.subscribe, topic) + subscribed.add(topic) + await self.server.wait_for_subscriptions(subscribed, self.latency) + await asyncio.wait_for(all_received.wait(), 5) + + self.assertEqual(received, [f"S{i}" for i in range(n)]) + subscribes = [m for m in self.server.inbox if m.get("action") == "subscribe"] + self.assertEqual(len(subscribes), 1 + n // 30) + self.assertEqual(errors, []) + await c.close() + await asyncio.wait_for(task, 5) + + async def test_cancel_while_connected(self) -> None: + received = asyncio.Event() + + async def processor(msgs): + received.set() + + c = self.client(subscriptions=["T.AAPL"]) + task = asyncio.create_task(c.connect(processor)) + await self.server.wait_for_subscriptions({"T.AAPL"}) + ws = c.websocket + assert ws is not None + self.addAsyncCleanup(ws.close) + # Cancel while the client waits for a message, not while one arrives: + # on Python 3.10 and 3.11, asyncio.wait_for() ignores a cancellation + # that coincides with the completion of what it waits for + # (https://bugs.python.org/issue42130). The acknowledgements are read + # before this message, so none is in flight once it is processed. + await self.server.send_data( + [{"ev": "T", "sym": "AAPL", "x": 1, "p": 1, "s": 1, "t": 1}] + ) + await asyncio.wait_for(received.wait(), 5) + + task.cancel() + done, _ = await asyncio.wait({task}, timeout=2) + self.assertIn(task, done, "connect() did not finish after cancel()") + self.assertTrue(task.cancelled()) + c.subscribe("T.MSFT") # does not raise after connect() ended + + +class SenderErrorTest(SubscriptionsTest): + async def test_send_error_is_raised_by_connect(self) -> None: + class Json: + loads = staticmethod(json.loads) + + @staticmethod + def dumps(obj: Any) -> str: + if obj.get("params") == "T.BAD": + raise TypeError("cannot encode") + return json.dumps(obj) + + c = self.client(subscriptions=["T.A"], custom_json=Json) + task = asyncio.create_task(c.connect(ignore)) + await self.server.wait_for_subscriptions({"T.A"}) + ws = c.websocket + assert ws is not None + self.addAsyncCleanup(ws.close) + c.subscribe("T.BAD") + with self.assertRaises(TypeError): + await asyncio.wait_for(task, 5) + + if __name__ == "__main__": unittest.main() From d611483ffd29bb51f9b7f8f506eecf2bc4c155d5 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:28:04 -0600 Subject: [PATCH 08/11] fix: match WebSocket wildcard collapse on the exact channel subscribe("T.*") and unsubscribe("T.*") discard the other subscribed topics on the channel, but they matched with t.startswith(channel), so subscribe("A.*") also dropped AM.AAPL and TX topics went with T.*. The test also matched the wildcard itself, so subscribing again to an already subscribed T.* unsubscribed T.*. The collapse now matches the channel exactly and keeps the wildcard being subscribed. As before, it applies only to subscribed topics, and topics are stored as given. --- massive/websocket/__init__.py | 27 ++++++++++++------- test_websocket/test_subscriptions.py | 39 ++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 9 deletions(-) diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 18bcf491..416c21d3 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -385,25 +385,32 @@ def _parse_subscriptions( parsed.append((s, topic, sym)) return parsed + def _collapse(self, scheduled: Set[str], channel: str, keep: Optional[str] = None): + """Unschedule the subscribed topics on exactly ``channel`` but ``keep``.""" + # subs is replaced, never changed in place, so iterating it is safe + for t in list(self.subs): + if t != keep and t.split(".", 1)[0].strip() == channel: + scheduled.discard(t) + def subscribe(self, *subscriptions: str): """ Subscribe to given subscriptions. Safe to call from any thread, before or while ``connect()`` runs. + Subscribing to a wildcard such as ``T.*`` replaces the other topics + subscribed on that channel, such as ``T.AAPL`` (but not ``TX.AAPL``). + :param subscriptions: Subscriptions (args) """ parsed = self._parse_subscriptions(subscriptions) with self._subs_lock: scheduled = set(self.scheduled_subs) - for s, topic, sym in parsed: + for s, channel, sym in parsed: scheduled.add(s) # If user subs to X.*, remove other X.\w+ if sym == "*": - for t in list(self.subs): - if t.startswith(topic): - scheduled.discard(t) - + self._collapse(scheduled, channel, keep=s) self.scheduled_subs = scheduled wake = self._request_resub() if wake: @@ -417,19 +424,21 @@ def unsubscribe(self, *subscriptions: str): Safe to call from any thread, before or while ``connect()`` runs. + Unsubscribing from a wildcard such as ``T.*`` also unsubscribes the + other topics subscribed on that channel, such as ``T.AAPL`` (but not + ``TX.AAPL``). + :param subscriptions: Subscriptions (args) """ parsed = self._parse_subscriptions(subscriptions) with self._subs_lock: scheduled = set(self.scheduled_subs) - for s, topic, sym in parsed: + for s, channel, sym in parsed: scheduled.discard(s) # If user unsubs to X.*, remove other X.\w+ if sym == "*": - for t in list(self.subs): - if t.startswith(topic): - scheduled.discard(t) + self._collapse(scheduled, channel) self.scheduled_subs = scheduled wake = self._request_resub() diff --git a/test_websocket/test_subscriptions.py b/test_websocket/test_subscriptions.py index 4d4d3012..27ba1273 100644 --- a/test_websocket/test_subscriptions.py +++ b/test_websocket/test_subscriptions.py @@ -299,5 +299,44 @@ def dumps(obj: Any) -> str: await asyncio.wait_for(task, 5) +class WildcardTest(SubscriptionsTest): + async def test_subscribe_collapses_only_the_same_channel(self) -> None: + c = self.client(subscriptions=["A.X", "T.AAPL", "TX.AAPL"]) + await self.start(c) + await self.server.wait_for_subscriptions({"A.X", "T.AAPL", "TX.AAPL"}) + c.subscribe("AM.*") # "A" is a prefix of "AM", not the same channel + await self.server.wait_for_subscriptions({"A.X", "T.AAPL", "TX.AAPL", "AM.*"}) + c.subscribe("T.*") + await self.server.wait_for_subscriptions({"A.X", "TX.AAPL", "AM.*", "T.*"}) + + async def test_collapse_leaves_topics_not_yet_subscribed(self) -> None: + # as before: only subscribed topics collapse + c = self.client(subscriptions=["T.AAPL"]) + c.subscribe("T.*") + self.assertEqual(c.scheduled_subs, {"T.AAPL", "T.*"}) + + async def test_unsubscribe_collapses_only_the_same_channel(self) -> None: + # T.AAPL after T.* (subscribing to T.AAPL first would collapse it) + c = self.client(subscriptions=["T.*", "T.AAPL", "TX.AAPL", "AM.X"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.AAPL", "T.*", "TX.AAPL", "AM.X"}) + c.unsubscribe("T.*", "A.*") + await self.server.wait_for_subscriptions({"TX.AAPL", "AM.X"}) + + async def test_subscribing_to_a_subscribed_wildcard_keeps_it(self) -> None: + c = self.client(subscriptions=["T.*"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.*"}) + c.subscribe("T.*", "Q.Z") + await self.server.wait_for_subscriptions({"T.*", "Q.Z"}) + + async def test_topics_are_kept_as_given(self) -> None: + # as before: the constructor does not collapse, and nothing is stripped + c = self.client(subscriptions=["T.AAPL", "T.*", " Q.X "]) + self.assertEqual(c.scheduled_subs, {"T.AAPL", "T.*", " Q.X "}) + c.unsubscribe(" Q.X ") + self.assertEqual(c.scheduled_subs, {"T.AAPL", "T.*"}) + + if __name__ == "__main__": unittest.main() From adda324276a95e334e5df7ab5030b057c1cc5ebf Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:30:51 -0600 Subject: [PATCH 09/11] feat: add WebSocketClient.set_subscriptions(desired) There was no way to replace the subscriptions in one step: a caller that knows the full set it wants had to compute the difference itself and call unsubscribe() and subscribe(), which another thread or the event loop could interleave with. set_subscriptions() atomically replaces the scheduled set; the client then sends at most one unsubscribe message for the removed topics followed by at most one subscribe message for the added ones. Removals go first so that a client at the server's subscription limit can swap topics; an update that also carries subscribe() or unsubscribe() changes made since the last one sends removals first too. Updates made only through subscribe() and unsubscribe() keep their order, additions first. The method is new, so no existing behaviour changes. --- docs/source/WebSocket.rst | 5 ++ massive/websocket/__init__.py | 45 +++++++++++++- test_websocket/test_subscriptions.py | 93 +++++++++++++++++++++++++++- 3 files changed, 140 insertions(+), 3 deletions(-) diff --git a/docs/source/WebSocket.rst b/docs/source/WebSocket.rst index c624d9d2..a3feb49e 100644 --- a/docs/source/WebSocket.rst +++ b/docs/source/WebSocket.rst @@ -34,6 +34,11 @@ Unsubscribe ============================ .. automethod:: massive.WebSocketClient.unsubscribe +============================ +Set subscriptions +============================ +.. automethod:: massive.WebSocketClient.set_subscriptions + ============================ Close ============================ diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 416c21d3..71727911 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -99,6 +99,9 @@ def __init__( # friends change on the caller's thread while connect() reads them on # the event loop's thread self._subs_lock = threading.Lock() + # set_subscriptions() was called since the last round: send removals + # first; guarded by _subs_lock + self._removals_first = False # while a connection is open: its loop, and the event that wakes its # sender task (see _start_sender) self._wake: Optional[Tuple[asyncio.AbstractEventLoop, asyncio.Event]] = None @@ -352,14 +355,23 @@ async def _reconcile(self, ws: ClientConnection): if not self.schedule_resub: return desired = set(self.scheduled_subs) + removals_first = self._removals_first + self._removals_first = False self.schedule_resub = False logger.debug("reconciling: %s %s", self.subs, desired) + added = desired.difference(self.subs) + removed = self.subs.difference(desired) try: - await self._subscribe(desired.difference(self.subs), ws) - await self._unsubscribe(self.subs.difference(desired), ws) + if removals_first: + await self._unsubscribe(removed, ws) + await self._subscribe(added, ws) + else: + await self._subscribe(added, ws) + await self._unsubscribe(removed, ws) except BaseException: # e.g. the connection closed: try again on the next connection with self._subs_lock: + self._removals_first |= removals_first self.schedule_resub = True raise self.subs = desired @@ -447,6 +459,35 @@ def unsubscribe(self, *subscriptions: str): for s, _, _ in parsed: logger.debug("sub undesired: %s", s) + def set_subscriptions(self, desired: Iterable[str]): + """ + Replace the subscriptions with ``desired``, in one step. + + Safe to call from any thread, before or while ``connect()`` runs. The + client then sends at most one unsubscribe message, for the topics no + longer desired, followed by at most one subscribe message, for the new + ones. Topics in both the old and the new set are left alone. Wildcards + such as ``T.*`` are ordinary topics here: nothing collapses. + + Unlike after ``subscribe()`` and ``unsubscribe()``, removals go first, + so a client at the server's subscription limit can swap topics; if + those methods were also called since the last update was sent, the + combined update sends removals first too. As a result, data that a + removed and an added topic both cover (e.g. ``T.AAPL`` replaced by + ``T.*``) is not delivered between the two messages. + + :param desired: Every topic to be subscribed to afterwards, e.g. + ``["T.AAPL", "Q.*"]``. Malformed topics are logged and skipped. + """ + topics = {s for s, _, _ in self._parse_subscriptions(desired)} + with self._subs_lock: + self.scheduled_subs = topics + self._removals_first = True + wake = self._request_resub() + if wake: + self._notify() + logger.debug("subs desired: %s", topics) + def unsubscribe_all(self): """ Unsubscribe from all subscriptions. diff --git a/test_websocket/test_subscriptions.py b/test_websocket/test_subscriptions.py index 27ba1273..51d1cfd1 100644 --- a/test_websocket/test_subscriptions.py +++ b/test_websocket/test_subscriptions.py @@ -3,7 +3,7 @@ import threading import unittest from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Set +from typing import Any, Dict, List, Set, Tuple from unittest import mock from fake_server import FakeServer @@ -338,5 +338,96 @@ async def test_topics_are_kept_as_given(self) -> None: self.assertEqual(c.scheduled_subs, {"T.AAPL", "T.*"}) +class FramesTest(SubscriptionsTest): + """Tests that check the exact subscribe and unsubscribe frames sent.""" + + def frames_after(self, n: int) -> List[Tuple[str, Set[str]]]: + """The frames the server received after its first ``n`` messages.""" + return [ + (m["action"], set(m["params"].split(","))) + for m in self.server.connection.inbox[n:] + ] + + async def wait_for_frame(self, action: str, topics: Set[str]): + await self.server.wait_for_message( + lambda m: m.get("action") == action + and set(m["params"].split(",")) == topics + ) + + +class SetSubscriptionsTest(FramesTest): + async def test_one_frame_per_action(self) -> None: + c = self.client(subscriptions=["T.A", "T.B", "T.C", "Q.*"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.A", "T.B", "T.C", "Q.*"}) + n = len(self.server.connection.inbox) + c.set_subscriptions(["T.B", "T.C", "T.D", "T.E"]) + self.assertEqual(c.scheduled_subs, {"T.B", "T.C", "T.D", "T.E"}) + await self.wait_for_frame("subscribe", {"T.D", "T.E"}) + self.assertEqual( + self.frames_after(n), + [("unsubscribe", {"T.A", "Q.*"}), ("subscribe", {"T.D", "T.E"})], + ) + await self.server.wait_for_subscriptions({"T.B", "T.C", "T.D", "T.E"}) + + async def test_subscribe_and_unsubscribe_send_additions_first(self) -> None: + # as before + c = self.client(subscriptions=["T.A"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.A"}) + n = len(self.server.connection.inbox) + c.subscribe("T.B") + c.unsubscribe("T.A") + await self.wait_for_frame("unsubscribe", {"T.A"}) + self.assertEqual( + self.frames_after(n), [("subscribe", {"T.B"}), ("unsubscribe", {"T.A"})] + ) + + async def test_combined_update_sends_removals_first(self) -> None: + c = self.client(subscriptions=["T.A", "T.B"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.A", "T.B"}) + n = len(self.server.connection.inbox) + c.subscribe("T.C") + c.set_subscriptions(["T.B", "T.C", "T.D"]) + await self.wait_for_frame("subscribe", {"T.C", "T.D"}) + self.assertEqual( + self.frames_after(n), + [("unsubscribe", {"T.A"}), ("subscribe", {"T.C", "T.D"})], + ) + + async def test_unchanged_set_sends_nothing(self) -> None: + c = self.client(subscriptions=["T.A"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.A"}) + n = len(self.server.connection.inbox) + c.set_subscriptions({"T.A"}) + # a change of its own, to show nothing was sent before it + c.subscribe("T.Z") + await self.wait_for_frame("subscribe", {"T.Z"}) + self.assertEqual(self.frames_after(n), [("subscribe", {"T.Z"})]) + + async def test_wildcards_do_not_collapse(self) -> None: + c = self.client(subscriptions=["T.A"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.A"}) + c.set_subscriptions(["T.A", "T.*"]) + await self.server.wait_for_subscriptions({"T.A", "T.*"}) + + async def test_before_connect(self) -> None: + c = self.client(subscriptions=["T.A"]) + c.set_subscriptions(["T.B", "T.*"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.B", "T.*"}) + self.assertEqual(self.frames_after(1), [("subscribe", {"T.B", "T.*"})]) + + async def test_malformed_topics_are_skipped(self) -> None: + c = self.client(subscriptions=["T.Z"]) + with mock.patch.object(websocket_module.logger, "warning") as warning: + c.set_subscriptions(["T.A", "bad", "T.B"]) + self.assertEqual(c.scheduled_subs, {"T.A", "T.B"}) + self.assertEqual([call.args[1] for call in warning.call_args_list], ["bad"]) + + if __name__ == "__main__": unittest.main() From 56ab8e656323c8263775472205f44e174b9839ee Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:29:39 -0600 Subject: [PATCH 10/11] feat: add WebSocketClient.resubscribe(*topics) The client sends a subscribe message only for topics missing from subs, so there was no way to send one again for a topic it already counts as subscribed, e.g. when the server never acknowledged it. resubscribe() adds such topics to the next subscribe message, together with any topics being added, and raises ValueError for a topic that is not subscribed or scheduled. The method is new, so no existing behaviour changes. --- docs/source/WebSocket.rst | 5 +++ massive/websocket/__init__.py | 33 +++++++++++++++- test_websocket/test_subscriptions.py | 59 ++++++++++++++++++++++++++++ 3 files changed, 96 insertions(+), 1 deletion(-) diff --git a/docs/source/WebSocket.rst b/docs/source/WebSocket.rst index a3feb49e..f557525b 100644 --- a/docs/source/WebSocket.rst +++ b/docs/source/WebSocket.rst @@ -39,6 +39,11 @@ Set subscriptions ============================ .. automethod:: massive.WebSocketClient.set_subscriptions +============================ +Resubscribe +============================ +.. automethod:: massive.WebSocketClient.resubscribe + ============================ Close ============================ diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 71727911..7f830586 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -102,6 +102,8 @@ def __init__( # set_subscriptions() was called since the last round: send removals # first; guarded by _subs_lock self._removals_first = False + # topics resubscribe() asked to send again; guarded by _subs_lock + self._resend: Set[str] = set() # while a connection is open: its loop, and the event that wakes its # sender task (see _start_sender) self._wake: Optional[Tuple[asyncio.AbstractEventLoop, asyncio.Event]] = None @@ -357,9 +359,11 @@ async def _reconcile(self, ws: ClientConnection): desired = set(self.scheduled_subs) removals_first = self._removals_first self._removals_first = False + resend = self._resend + self._resend = set() self.schedule_resub = False logger.debug("reconciling: %s %s", self.subs, desired) - added = desired.difference(self.subs) + added = desired.difference(self.subs) | resend.intersection(desired) removed = self.subs.difference(desired) try: if removals_first: @@ -372,6 +376,7 @@ async def _reconcile(self, ws: ClientConnection): # e.g. the connection closed: try again on the next connection with self._subs_lock: self._removals_first |= removals_first + self._resend |= resend self.schedule_resub = True raise self.subs = desired @@ -488,6 +493,32 @@ def set_subscriptions(self, desired: Iterable[str]): self._notify() logger.debug("subs desired: %s", topics) + def resubscribe(self, *topics: str): + """ + Send a subscribe message again for topics the client already counts + as subscribed, e.g. ones the server never acknowledged. + + Safe to call from any thread, before or while ``connect()`` runs. The + topics go into the next subscribe message, together with any topics + being added. A topic unsubscribed before that message is sent is not + sent. + + :param topics: Topics to send again (args). Malformed topics are + logged and skipped. + :raises ValueError: If a well-formed topic is not subscribed (or + scheduled to be); nothing is then sent. + """ + again = {s for s, _, _ in self._parse_subscriptions(topics)} + with self._subs_lock: + unknown = again.difference(self.scheduled_subs) + if unknown: + raise ValueError(f"not subscribed: {sorted(unknown)}") + self._resend |= again + wake = self._request_resub() + if wake: + self._notify() + logger.debug("subs to send again: %s", again) + def unsubscribe_all(self): """ Unsubscribe from all subscriptions. diff --git a/test_websocket/test_subscriptions.py b/test_websocket/test_subscriptions.py index 51d1cfd1..117d8c80 100644 --- a/test_websocket/test_subscriptions.py +++ b/test_websocket/test_subscriptions.py @@ -429,5 +429,64 @@ async def test_malformed_topics_are_skipped(self) -> None: self.assertEqual([call.args[1] for call in warning.call_args_list], ["bad"]) +class ResubscribeTest(FramesTest): + async def connected(self, *topics: str) -> WebSocketClient: + """A client subscribed to ``topics``, acknowledged by the server.""" + c = self.client(subscriptions=list(topics)) + await self.start(c) + await self.server.wait_for_subscriptions(set(topics)) + return c + + async def test_sends_an_acknowledged_topic_again(self) -> None: + c = await self.connected("T.A", "T.B") + n = len(self.server.connection.inbox) + c.resubscribe("T.A") + await self.wait_for_frame("subscribe", {"T.A"}) + # a change of its own, to show nothing else was sent in between + c.subscribe("T.Z") + await self.wait_for_frame("subscribe", {"T.Z"}) + self.assertEqual( + self.frames_after(n), [("subscribe", {"T.A"}), ("subscribe", {"T.Z"})] + ) + + async def test_joins_the_subscribe_message_for_new_topics(self) -> None: + c = await self.connected("T.A", "T.B") + n = len(self.server.connection.inbox) + c.set_subscriptions(["T.A", "T.B", "T.C"]) + c.resubscribe("T.A") + await self.wait_for_frame("subscribe", {"T.A", "T.C"}) + self.assertEqual(self.frames_after(n), [("subscribe", {"T.A", "T.C"})]) + + async def test_not_sent_if_unsubscribed_first(self) -> None: + c = await self.connected("T.A", "T.B") + n = len(self.server.connection.inbox) + c.resubscribe("T.A") + c.unsubscribe("T.A") + await self.wait_for_frame("unsubscribe", {"T.A"}) + c.subscribe("T.Z") + await self.wait_for_frame("subscribe", {"T.Z"}) + self.assertEqual( + self.frames_after(n), [("unsubscribe", {"T.A"}), ("subscribe", {"T.Z"})] + ) + + async def test_topic_must_be_subscribed(self) -> None: + c = await self.connected("T.A") + n = len(self.server.connection.inbox) + with self.assertRaises(ValueError): + c.resubscribe("T.A", "T.B") + c.subscribe("T.Z") + await self.wait_for_frame("subscribe", {"T.Z"}) + self.assertEqual(self.frames_after(n), [("subscribe", {"T.Z"})]) + + async def test_malformed_topics_are_skipped(self) -> None: + c = await self.connected("T.A", "T.B") + n = len(self.server.connection.inbox) + with mock.patch.object(websocket_module.logger, "warning") as warning: + c.resubscribe("bad", "T.A") + self.assertEqual([call.args[1] for call in warning.call_args_list], ["bad"]) + await self.wait_for_frame("subscribe", {"T.A"}) + self.assertEqual(self.frames_after(n), [("subscribe", {"T.A"})]) + + if __name__ == "__main__": unittest.main() From adfb99bd26eb114e03f598dc9e7ff2659cb9192d Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:32:49 -0600 Subject: [PATCH 11/11] fix: keep scheduled WebSocket subscriptions across reconnects After an abnormal close, connect() replaced scheduled_subs with subs, the topics the old connection had applied. A subscribe() or unsubscribe() made before the close but not yet sent, or made during the outage, was lost, and an unsubscribed topic could come back. scheduled_subs is now the authoritative desired set and is left alone: only subs is reset, and the new connection subscribes to the scheduled set. Likewise, each connect() call starts with empty subs, so calling it again after close() or after cancelling it subscribes to the scheduled set instead of assuming the old connection's topics are still subscribed. --- massive/websocket/__init__.py | 9 +++- test_websocket/test_subscriptions.py | 61 ++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 7f830586..af1964bd 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -132,6 +132,11 @@ async def connect( :param close_timeout: How long to wait for handshake when calling .close. :raises AuthError: If invalid API key is supplied. """ + # a new connection has no subscriptions yet, whatever an earlier + # connect() call left in subs + with self._subs_lock: + self.subs = set() + self.schedule_resub = True reconnects = 0 logger.debug("connect: %s", self.url) ssl_context = self._ssl_context() @@ -205,8 +210,10 @@ async def connect( logger.debug("connection closed (ERR): %s", e) self._flush_server_errors(force=True) reconnects += 1 + # scheduled_subs is the desired set: keep changes made before + # or during the outage, and subscribe to all of it again on + # the new connection with self._subs_lock: - self.scheduled_subs = set(self.subs) self.subs = set() self.schedule_resub = True if self.max_reconnects is not None and reconnects > self.max_reconnects: diff --git a/test_websocket/test_subscriptions.py b/test_websocket/test_subscriptions.py index 117d8c80..64f49b8e 100644 --- a/test_websocket/test_subscriptions.py +++ b/test_websocket/test_subscriptions.py @@ -488,5 +488,66 @@ async def test_malformed_topics_are_skipped(self) -> None: self.assertEqual(self.frames_after(n), [("subscribe", {"T.A"})]) +class ReconnectTest(SubscriptionsTest): + async def test_changes_during_an_outage_are_kept(self) -> None: + c = self.client(subscriptions=["T.A", "T.B"]) + await self.start(c) + await self.server.wait_for_subscriptions({"T.A", "T.B"}) + + # drop the connection, and hold the next one before authentication, + # so the client is between connections while the changes are made + self.server.script_auth() + await self.server.close_connections(code=1011) + await self.server.wait_for( + lambda: len(self.server.connections) == 2 + and self.server.connection.api_key is not None + ) + c.subscribe("T.C") + c.unsubscribe("T.A") + await self.server.connection.close(code=1011) + + await self.server.wait_for_connections(2) + self.assertEqual(len(self.server.connections), 3) + await self.server.wait_for_subscriptions({"T.B", "T.C"}) + self.assertEqual(c.scheduled_subs, {"T.B", "T.C"}) + + async def test_connect_again_after_close(self) -> None: + c = self.client(subscriptions=["T.A"]) + task = await self.start(c) + await self.server.wait_for_subscriptions({"T.A"}) + await c.close() + await asyncio.wait_for(task, 5) + + c.subscribe("T.B") + await self.start(c) + await self.server.wait_for_subscriptions({"T.A", "T.B"}) + + async def test_connect_again_after_cancel(self) -> None: + received = asyncio.Event() + + async def processor(msgs): + received.set() + + c = self.client(subscriptions=["T.A"]) + task = asyncio.create_task(c.connect(processor)) + await self.server.wait_for_subscriptions({"T.A"}) + ws = c.websocket + assert ws is not None + self.addAsyncCleanup(ws.close) + # cancel once idle; see WakeUpTest.test_cancel_while_connected + await self.server.send_data( + [{"ev": "T", "sym": "A", "x": 1, "p": 1, "s": 1, "t": 1}] + ) + await asyncio.wait_for(received.wait(), 5) + task.cancel() + done, _ = await asyncio.wait({task}, timeout=2) + self.assertIn(task, done) + + await self.start(c) + await self.server.wait_for_subscriptions({"T.A"}) + c.subscribe("T.B") + await self.server.wait_for_subscriptions({"T.A", "T.B"}) + + if __name__ == "__main__": unittest.main()