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 diff --git a/docs/source/WebSocket.rst b/docs/source/WebSocket.rst index c624d9d2..f557525b 100644 --- a/docs/source/WebSocket.rst +++ b/docs/source/WebSocket.rst @@ -34,6 +34,16 @@ Unsubscribe ============================ .. automethod:: massive.WebSocketClient.unsubscribe +============================ +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 7463b4cb..af1964bd 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -1,15 +1,30 @@ 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 -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 @@ -76,8 +91,22 @@ 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() + # 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 self._server_errors: Counter = Counter() self._server_errors_since = 0.0 self.ca_certs = ca_certs @@ -103,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() @@ -132,39 +166,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: - 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 - - 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 @@ -175,9 +210,12 @@ 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 + # 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.subs = set() + self.schedule_resub = True if self.max_reconnects is not None and reconnects > self.max_reconnects: break continue @@ -242,23 +280,113 @@ 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})) + + 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: + if not self.schedule_resub: + return + 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) | resend.intersection(desired) + removed = self.subs.difference(desired) + try: + 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._resend |= resend + self.schedule_resub = True + raise + self.subs = desired @staticmethod def _parse_subscription(s: str): @@ -270,53 +398,145 @@ 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 _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) """ - 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, channel, sym in parsed: + scheduled.add(s) + # If user subs to X.*, remove other X.\w+ + if sym == "*": + self._collapse(scheduled, channel, keep=s) + self.scheduled_subs = scheduled + wake = self._request_resub() + if wake: + self._notify() + 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. + + 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) """ - 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, channel, sym in parsed: + scheduled.discard(s) + + # If user unsubs to X.*, remove other X.\w+ + if sym == "*": + self._collapse(scheduled, channel) + + self.scheduled_subs = scheduled + wake = self._request_resub() + if wake: + self._notify() + 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 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 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. + + 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() + wake = self._request_resub() + if wake: + self._notify() async def close(self): """ 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] 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_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() 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( 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() diff --git a/test_websocket/test_subscriptions.py b/test_websocket/test_subscriptions.py new file mode 100644 index 00000000..64f49b8e --- /dev/null +++ b/test_websocket/test_subscriptions.py @@ -0,0 +1,553 @@ +import asyncio +import json +import threading +import unittest +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Set, Tuple +from unittest import mock + +from fake_server import FakeServer + +from massive import WebSocketClient +from massive import websocket as websocket_module + + +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) + + +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) + + +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.*"}) + + +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"]) + + +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"})]) + + +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()