diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index a1ae9ada..e49a8035 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -14,7 +14,7 @@ jobs: fail-fast: false matrix: os: [ubuntu-latest, windows-latest] - python-version: ['3.10', '3.11'] + python-version: ['3.10', '3.11', '3.12', '3.13', '3.14'] runs-on: ${{ matrix.os }} name: ${{ matrix.os }} Unit test ${{ matrix.python-version }} steps: @@ -32,3 +32,32 @@ jobs: run: poetry install - name: Unit tests run: make test + test-free-threaded: + strategy: + fail-fast: false + matrix: + python-version: ['3.13t', '3.14t'] + runs-on: ubuntu-latest + name: ubuntu-latest Unit test ${{ matrix.python-version }} (PYTHON_GIL=0) + steps: + - uses: actions/checkout@v3 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Setup Poetry + # pipx is Poetry's recommended installation method, and it keeps Poetry + # on the runner's regular Python, independent of the interpreter under test + run: pipx install poetry + - name: Install pypi deps + # the dev group (linters, docs, orjson for the examples) is not needed + # to run the tests, and orjson does not support free-threaded Python + run: | + poetry env use python${{ matrix.python-version }} + poetry install --only main,test + - name: Unit tests + env: + PYTHON_GIL: '0' + run: | + poetry run python -c 'import sys; assert not sys._is_gil_enabled(), sys.version' + make test diff --git a/docs/source/WebSocket.rst b/docs/source/WebSocket.rst index c624d9d2..e4a6b412 100644 --- a/docs/source/WebSocket.rst +++ b/docs/source/WebSocket.rst @@ -39,6 +39,11 @@ Close ============================ .. automethod:: massive.WebSocketClient.close +============================ +Request close +============================ +.. automethod:: massive.WebSocketClient.request_close + .. _Stocks getting started: https://massive.com/docs/stocks/ws_getting-started .. _Options getting started: https://massive.com/docs/options/ws_getting-started .. _Forex getting started: https://massive.com/docs/forex/ws_getting-started diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 7463b4cb..3dbb2b74 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -5,11 +5,18 @@ import logging import json import asyncio +import random 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, + InvalidMessage, + InvalidStatus, +) from ..certs import ca_bundle_path from ..logging import get_logger import logging @@ -32,10 +39,12 @@ def __init__( raw: bool = False, verbose: bool = False, subscriptions: Optional[List[str]] = None, - max_reconnects: Optional[int] = 5, + max_reconnects: Optional[int] = None, secure: bool = True, custom_json: Optional[Any] = None, ca_certs: Optional[str] = None, + reconnect_backoff: float = 15.0, + max_reconnect_backoff: float = 120.0, **kwargs, ): """ @@ -46,11 +55,31 @@ def __init__( :param raw: Whether to pass raw Union[str, bytes] to user callback. :param verbose: Whether to log client and server status messages. :param subscriptions: List of subscription parameters. - :param max_reconnects: How many times to reconnect on network outage before ending .connect event loop. + :param max_reconnects: How many consecutive reconnect attempts may fail + (the connection is refused, or ends before or after authenticating) + before .connect raises the last error. The count starts over whenever + a connection authenticates. None, the default, retries forever. Only + transient failures are retried and the wait between attempts grows, + while other errors (e.g. HTTP 4xx) are raised at once, so with the + default .run is safe in a long-lived process. :param custom_json: Optional module exposing loads/dumps functions (similar to Python's json module) to be used for JSON conversions. :param ca_certs: Path to a CA bundle to verify TLS with. Defaults to the SSL_CERT_FILE or REQUESTS_CA_BUNDLE environment variable, then certifi. + :param reconnect_backoff: Seconds to wait before reconnecting after the + connection ends or an attempt to connect fails. The wait doubles with + each consecutive failure and gets up to 20% random jitter, but never + leaves [reconnect_backoff, max_reconnect_backoff]. Defaults to 15, the + minimum Massive asks clients to wait; 0 disables the wait. + :param max_reconnect_backoff: The longest wait between reconnect + attempts, jitter included. + :raises ValueError: If reconnect_backoff is negative or greater than + max_reconnect_backoff. :return: A client. """ + if reconnect_backoff < 0 or max_reconnect_backoff < reconnect_backoff: + raise ValueError( + "need 0 <= reconnect_backoff <= max_reconnect_backoff, got " + f"{reconnect_backoff} and {max_reconnect_backoff}" + ) if api_key is None: raise AuthError( f"Must specify env var {env_key} or pass api_key in constructor" @@ -85,6 +114,15 @@ def __init__( self.json = custom_json else: self.json = json + # set while connect() runs; close() and request_close() stop it + self._loop: Optional[asyncio.AbstractEventLoop] = None + self._close_requested = threading.Event() + self._close_event: Optional[asyncio.Event] = None + self._close_task: Optional[asyncio.Future] = None + # whether the server said max_connections during the last handshake + self.max_connections_reached = False + self.reconnect_backoff = reconnect_backoff + self.max_reconnect_backoff = max_reconnect_backoff # https://websockets.readthedocs.io/en/stable/reference/client.html#opening-a-connection async def connect( @@ -99,27 +137,98 @@ async def connect( """ Connect to websocket server and run `processor(msg)` on every new `msg`. + When the connection ends, whether the server closed it cleanly or not, + or an attempt to connect fails for a reason that may pass, the client + waits (see `reconnect_backoff`), reconnects, and subscribes again to + every scheduled subscription. Such reasons are network errors (refused, + unreachable, timed out, a DNS lookup failure, or a TLS error other + than certificate verification), the connection closing during the + handshake, and HTTP 408, 429 or 5xx. Other errors, e.g. any other HTTP + 4xx or a TLS certificate that cannot be verified, are raised at once. + It returns once `close()` or `request_close()` is called. + + If the server sends a `max_connections` status during the handshake + (the account has too many connections open), `max_connections_reached` + is set until a later handshake succeeds, and in raw mode the message is + also passed to `processor`. The server then closes the connection, and + the client reconnects after the backoff. + :param processor: The callback to process messages. :param close_timeout: How long to wait for handshake when calling .close. + :param kwargs: Passed to websockets' connect() for each attempt. :raises AuthError: If invalid API key is supplied. + :raises TypeError: If `process_exception` is passed: connect() decides + itself which errors to retry. + :raises Exception: If more than `max_reconnects` consecutive attempts + fail, the error that ended the last one, e.g. ConnectionClosedOK if + the server closed the last connection cleanly. """ - reconnects = 0 + self._loop = asyncio.get_running_loop() + self._close_event = asyncio.Event() + try: + await self._connect(processor, close_timeout, **kwargs) + # it returned because close was requested, which is now done + self._close_requested.clear() + finally: + self._loop = None + self._close_event = None + closing, self._close_task = self._close_task, None + if closing is not None: + await asyncio.gather(closing, return_exceptions=True) + + async def _connect( + self, + processor: Union[ + Callable[[List[WebSocketMessage]], Awaitable], + Callable[[Union[str, bytes]], Awaitable], + ], + close_timeout: int, + **kwargs, + ): + # consecutive failed attempts, which set the backoff and are limited + # by max_reconnects + failures = 0 logger.debug("connect: %s", self.url) ssl_context = self._ssl_context() - last_exc = None - async for s in connect( - self.url, - close_timeout=close_timeout, - ssl=ssl_context, - process_exception=lambda exc: None, - **kwargs, - ): + if "process_exception" in kwargs: + raise TypeError( + "process_exception is not supported: connect() decides which " + "errors to retry" + ) + + def new_attempt() -> Awaitable[ClientConnection]: + return connect( + self.url, close_timeout=close_timeout, ssl=ssl_context, **kwargs + ) + + # build one up front so that invalid arguments raise at once + new_attempt() + last_exc: Optional[Exception] = None + delay = 0.0 + while not self._close_requested.is_set(): + if failures > 0 and await self._wait_for_close(delay): + return + try: + s = await self._open(new_attempt()) + except Exception as e: + if not self._is_transient(e): + raise + last_exc = e + failures += 1 + if self._too_many_failures(failures): + break + delay = self._reconnect_delay(failures - 1) + logger.warning("connect failed; reconnecting in %.1f s: %s", delay, e) + continue + if s is None: + return self.websocket = s try: msg = await s.recv() logger.debug("connected: %s", msg) - self._log_server_statuses(self.json.loads(msg)) + msg_parsed = self.json.loads(msg) + self._log_server_statuses(msg_parsed) logger.debug("authing...") await s.send( self.json.dumps({"action": "auth", "params": self.api_key}) @@ -129,6 +238,14 @@ async def connect( logger.debug("authed: %s", auth_msg) if auth_msg_parsed[0]["status"] == "auth_failed": raise AuthError(auth_msg_parsed[0]["message"]) + self.max_connections_reached = False + for frame, parsed in ((msg, msg_parsed), (auth_msg, auth_msg_parsed)): + if any(m.get("status") == "max_connections" for m in parsed): + self.max_connections_reached = True + if self.raw: + await processor(frame) # type: ignore + if not self.max_connections_reached: + failures = 0 # e.g. max_connections, which the server sends before closing self._log_server_statuses(auth_msg_parsed) self._flush_server_errors(force=True) @@ -165,30 +282,108 @@ async def connect( if len(cmsg) > 0: await processor(cmsg) # type: ignore - except ConnectionClosedOK as e: - logger.debug("connection closed (OK): %s", e) + except ConnectionClosed as e: + ok = isinstance(e, ConnectionClosedOK) + logger.debug("connection closed (%s): %s", "OK" if ok else "ERR", e) # the server sends a status (e.g. slow_consumer) before closing self._flush_server_errors(force=True) - return - except ConnectionClosedError as e: + if self._close_requested.is_set(): + return + # the server can close cleanly too, e.g. after max_connections last_exc = e - logger.debug("connection closed (ERR): %s", e) - self._flush_server_errors(force=True) - reconnects += 1 - self.scheduled_subs = set(self.subs) + failures += 1 + if self._too_many_failures(failures): + break + delay = self._reconnect_delay(failures - 1) + logger.warning( + "connection closed; reconnecting in %.1f s: %s", delay, e + ) + finally: + self.websocket = None + # scheduled_subs is the desired set, including changes not sent + # yet: keep it, and subscribe all of it on the next connection self.subs = set() self.schedule_resub = True - if self.max_reconnects is not None and reconnects > self.max_reconnects: - break - continue + await s.close() - if ( - last_exc - and self.max_reconnects is not None - and reconnects > self.max_reconnects - ): + if last_exc is not None and self._too_many_failures(failures): raise last_exc + def _too_many_failures(self, failures: int) -> bool: + return self.max_reconnects is not None and failures > self.max_reconnects + + @staticmethod + def _is_transient(exc: Exception) -> bool: + """Whether a failed attempt to connect is worth retrying.""" + if isinstance(exc, InvalidStatus): + status = exc.response.status_code + return status >= 500 or status in (408, 429) + if isinstance(exc, InvalidMessage): + # websockets 15: the connection closed or failed before a response + return isinstance(exc.__cause__, (EOFError, OSError)) + if isinstance(exc, ssl.SSLCertVerificationError): + return False + # e.g. refused, unreachable, timed out, DNS, or another TLS error; + # websockets 14 raises EOFError if the connection closes before a response + return isinstance(exc, (OSError, EOFError, asyncio.TimeoutError)) + + def _reconnect_delay(self, attempt: int) -> float: + """ + The wait before reconnect attempt `attempt`, counting from 0. + + The base doubles from `reconnect_backoff` up to `max_reconnect_backoff`; + jitter only lengthens it, by up to 20%, and is clamped to the cap. + """ + # bound the exponent so a long outage cannot overflow the float + base = min( + self.reconnect_backoff * 2.0 ** min(attempt, 1000), + self.max_reconnect_backoff, + ) + return random.uniform(base, min(1.2 * base, self.max_reconnect_backoff)) + + async def _open( + self, connector: Awaitable[ClientConnection] + ) -> Optional[ClientConnection]: + """Open a connection, or return None if close is requested first.""" + assert self._close_event is not None + attempt = asyncio.ensure_future(connector) + closing = asyncio.ensure_future(self._close_event.wait()) + try: + await asyncio.wait({attempt, closing}, return_when=asyncio.FIRST_COMPLETED) + except BaseException: + # e.g. connect() was cancelled: don't leave a connection open + closing.cancel() + await asyncio.shield(self._abandon(attempt)) + raise + stop = closing.done() or self._close_requested.is_set() + closing.cancel() + if not stop: + return attempt.result() + await self._abandon(attempt) + return None + + @staticmethod + async def _abandon(attempt: "asyncio.Future[ClientConnection]"): + """ + Cancel an attempt to connect, or close the connection it opened. + + Waits for the attempt to end, which websockets bounds by its + `open_timeout`. + """ + attempt.cancel() + await asyncio.wait({attempt}) + if not attempt.cancelled() and attempt.exception() is None: + await attempt.result().close() + + async def _wait_for_close(self, timeout: float) -> bool: + """Wait up to `timeout` seconds; return whether close was requested.""" + assert self._close_event is not None + try: + await asyncio.wait_for(self._close_event.wait(), timeout) + except asyncio.TimeoutError: + pass + return self._close_requested.is_set() + def _ssl_context(self) -> Optional[ssl.SSLContext]: # darwin needs some extra <3 if not self.url.startswith("wss://"): @@ -318,15 +513,51 @@ def unsubscribe_all(self): self.scheduled_subs = set() self.schedule_resub = True + def request_close(self): + """ + Close the websocket connection and make `connect()` return. + + Unlike `close()`, this can be called from any thread, and it does not + wait: `connect()` returns soon after, including while it is + connecting or waiting to reconnect. If no `connect()` is running, the + next one returns without connecting; the request stays pending until a + `connect()` returns because of it. + """ + logger.debug("close requested") + self._close_requested.set() + loop = self._loop + if loop is None: + return + try: + loop.call_soon_threadsafe(self._on_close_requested) + except RuntimeError: + # the loop is closed, so connect() has returned + pass + + def _on_close_requested(self): + # runs on the event loop of connect(); the request may be stale, e.g. + # queued for an earlier connect() on the same loop + if not self._close_requested.is_set(): + return + if self._close_event is not None: + self._close_event.set() + if self.websocket is not None: + self._close_task = asyncio.ensure_future(self.websocket.close()) + async def close(self): """ - Close the websocket connection. + Close the websocket connection and stop `connect()` from reconnecting. """ logger.debug("closing") + if self._loop is not None: + self._close_requested.set() + if self._close_event is not None: + self._close_event.set() self._flush_server_errors(force=True) if self.websocket: await self.websocket.close() self.websocket = None - else: + elif self._loop is None: + # (while connect() runs, it may be between connections) logger.warning("no websocket open to close") diff --git a/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/test_concurrency.py b/test_rest/test_concurrency.py new file mode 100644 index 00000000..1b1c4602 --- /dev/null +++ b/test_rest/test_concurrency.py @@ -0,0 +1,118 @@ +import json +import threading +import unittest +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import List, Tuple +from urllib.parse import parse_qsl, urlparse + +import pook # type: ignore + +from massive import RESTClient +from massive.rest.models.request import RequestOptionBuilder, X_MASSIVE_EDGE_ID + + +class EchoHandler(BaseHTTPRequestHandler): + # keep-alive, so the client reuses pooled connections across threads + protocol_version = "HTTP/1.1" + + def do_GET(self): + url = urlparse(self.path) + body = json.dumps( + { + "results": { + "path": url.path, + "query": dict(parse_qsl(url.query)), + "edge_id": self.headers.get(X_MASSIVE_EDGE_ID), + "authorization": self.headers.get("Authorization"), + } + } + ).encode() + with self.server.lock: # type: ignore[attr-defined] + self.server.count += 1 # type: ignore[attr-defined] + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format, *args): + pass + + +class ConcurrencyTest(unittest.TestCase): + """One RESTClient shared by many threads against a local HTTP server.""" + + threads = 16 + calls = 200 + + @classmethod + def setUpClass(cls): + # other suites leave pook intercepting urllib3; let requests through + cls.pook_was_active = pook.isactive() + pook.disable() + + @classmethod + def tearDownClass(cls): + if cls.pook_was_active: + pook.activate() + + def setUp(self): + httpd = ThreadingHTTPServer(("127.0.0.1", 0), EchoHandler) + httpd.daemon_threads = True + httpd.lock = threading.Lock() # type: ignore[attr-defined] + httpd.count = 0 # type: ignore[attr-defined] + thread = threading.Thread(target=httpd.serve_forever, daemon=True) + thread.start() + self.addCleanup(thread.join) + self.addCleanup(httpd.server_close) + self.addCleanup(httpd.shutdown) + self.httpd = httpd + host, port = httpd.server_address[:2] + self.base = f"http://{host}:{port}" + + def test_shared_client(self): + c = RESTClient( + "key", base=self.base, retries=0, maxsize=self.threads, pagination=False + ) + headers_before = dict(c.headers) + barrier = threading.Barrier(self.threads, timeout=10) + + def run(t: int) -> List[Tuple[dict, dict]]: + barrier.wait() + results = [] + for i in range(self.calls): + expected = { + "path": f"/v1/thread/{t}/call/{i}", + "query": {"t": str(t), "i": str(i)}, + "edge_id": f"edge-{t}-{i}" if i % 2 else None, + "authorization": "Bearer key", + } + options = ( + RequestOptionBuilder(edge_id=f"edge-{t}-{i}", edge_ip_address="ip") + if i % 2 + else None + ) + got = c._get( + path=expected["path"], + params={"t": t, "i": i}, + result_key="results", + options=options, + ) + results.append((expected, got)) + return results + + with ThreadPoolExecutor(self.threads) as pool: + futures = [pool.submit(run, t) for t in range(self.threads)] + results = [r for f in futures for r in f.result()] + + self.assertEqual(len(results), self.threads * self.calls) + for expected, got in results: + self.assertEqual(got, expected) + self.assertEqual(self.httpd.count, self.threads * self.calls) # type: ignore[attr-defined] + # per-request options must not leak into the shared default headers + self.assertEqual(c.headers, headers_before) + + +if __name__ == "__main__": + unittest.main() diff --git a/test_websocket/fake_server.py b/test_websocket/fake_server.py new file mode 100644 index 00000000..bf77a4b8 --- /dev/null +++ b/test_websocket/fake_server.py @@ -0,0 +1,342 @@ +""" +A scriptable fake Massive WebSocket server for tests. + +Unlike ``mock_server``, which listens on a fixed port and keeps module-level +state, each ``FakeServer`` listens on a free port and records everything per +instance, so tests can run several servers and assert on exactly what a client +sent. Run it inside the test's event loop:: + + async with FakeServer() as server: + c = WebSocketClient(api_key="key", feed=server.feed, secure=False) + c.subscribe("T.AAPL") + task = asyncio.create_task(c.connect(processor)) + await server.wait_for_subscriptions({"T.AAPL"}) + await server.send_data([{"ev": "T", "sym": "AAPL", "p": 1.0}]) + await server.close_connections(code=1011) # simulate an outage + +By default the server speaks the Massive protocol: it sends ``connected``, +answers ``auth`` with ``auth_success`` (or ``auth_failed`` for a key in +``reject_keys``), and acknowledges each subscribed or unsubscribed topic with a +``success`` status. Scripted replies and hooks override that behaviour. +""" + +import asyncio +import json +from http import HTTPStatus +from typing import ( + Any, + Awaitable, + Callable, + Dict, + Iterable, + List, + Optional, + Set, +) + +from websockets.asyncio.server import Server, ServerConnection, serve +from websockets.exceptions import ConnectionClosed +from websockets.http11 import Request, Response + + +def status(status: str, message: str) -> Dict[str, Any]: + """Build a Massive status message, e.g. ``status("max_connections", "...")``.""" + return {"ev": "status", "status": status, "message": message} + + +class AuthReply: + """ + A scripted reply to one ``auth`` message. + + :param frames: The messages to send, as a JSON array. + :param close_code: If set, close the connection with this code afterwards. + :param close_reason: The reason sent with the close frame. + """ + + def __init__( + self, + frames: List[Dict[str, Any]], + close_code: Optional[int] = None, + close_reason: str = "", + ): + self.frames = frames + self.close_code = close_code + self.close_reason = close_reason + + +class FakeConnection: + """One client connection accepted by a ``FakeServer``.""" + + def __init__(self, server: "FakeServer", websocket: ServerConnection, index: int): + self.server = server + self.websocket = websocket + self.index = index + self.path = websocket.request.path if websocket.request else "" + self.api_key: Optional[str] = None + self.authenticated = False + # every message the client sent, decoded from JSON + self.inbox: List[Any] = [] + # topics the client is subscribed to on this connection + self.subscriptions: Set[str] = set() + self.closed = asyncio.Event() + + @property + def open(self) -> bool: + return not self.closed.is_set() + + async def send(self, frames: List[Dict[str, Any]]): + """Send a JSON array of messages to this client.""" + await self.websocket.send(json.dumps(frames)) + + async def send_raw(self, message: str): + """Send a message as-is.""" + await self.websocket.send(message) + + async def close(self, code: int = 1000, reason: str = ""): + """Close with a close frame; 1000 is a clean close, e.g. 1011 is not.""" + await self.websocket.close(code, reason) + await self.closed.wait() + + async def drop(self): + """Drop the TCP connection without a close frame.""" + transport = self.websocket.transport + if transport is not None: + transport.abort() + await self.closed.wait() + + +class FakeServer: + """ + A fake Massive WebSocket server listening on ``host`` and a free port. + + :param host: The interface to listen on. + :param valid_keys: If set, only these API keys authenticate. + :param reject_keys: API keys answered with ``auth_failed``. + :param ack_subscriptions: Whether to acknowledge subscribe/unsubscribe. + :param reject_topics: Topics answered with an ``error`` status instead of + being subscribed. + :param on_message: Called with every decoded client message after it has + been recorded and handled; use it to script replies. + """ + + def __init__( + self, + host: str = "127.0.0.1", + valid_keys: Optional[Iterable[str]] = None, + reject_keys: Iterable[str] = (), + ack_subscriptions: bool = True, + reject_topics: Iterable[str] = (), + on_message: Optional[Callable[[FakeConnection, Any], Awaitable[None]]] = None, + ): + self.host = host + self.port = 0 + self.valid_keys = set(valid_keys) if valid_keys is not None else None + self.reject_keys = set(reject_keys) + self.ack_subscriptions = ack_subscriptions + self.reject_topics = set(reject_topics) + self.on_message = on_message + # while True, new connections get an HTTP 503 instead of a handshake + self.reject_handshakes = False + self.connections: List[FakeConnection] = [] + self._auth_replies: List[AuthReply] = [] + self._server: Optional[Server] = None + + # lifecycle + + async def start(self) -> "FakeServer": + self._server = await serve( + self._handler, + self.host, + self.port, + process_request=self._process_request, + ) + self.port = next(iter(self._server.sockets)).getsockname()[1] + return self + + async def stop(self): + if self._server is not None: + self._server.close() + await self._server.wait_closed() + self._server = None + + async def __aenter__(self) -> "FakeServer": + return await self.start() + + async def __aexit__(self, *exc): + await self.stop() + + @property + def feed(self) -> str: + """The ``feed`` to pass to ``WebSocketClient`` (with ``secure=False``).""" + return f"{self.host}:{self.port}" + + # scripting + + def script_auth( + self, + *frames: Dict[str, Any], + close_code: Optional[int] = None, + close_reason: str = "", + ): + """ + Queue the reply to the next ``auth`` message, replacing the default. + + Replies are used in order, one per connection, e.g. + ``script_auth(status("max_connections", "..."), close_code=1008)``. + """ + self._auth_replies.append(AuthReply(list(frames), close_code, close_reason)) + + @property + def connection(self) -> FakeConnection: + """The most recent connection.""" + if not self.connections: + raise AssertionError("no client has connected") + return self.connections[-1] + + @property + def open_connections(self) -> List[FakeConnection]: + return [c for c in self.connections if c.open] + + @property + def inbox(self) -> List[Any]: + """Every message received, from all connections, in order per connection.""" + return [m for c in self.connections for m in c.inbox] + + @property + def subscriptions(self) -> Set[str]: + """The topics subscribed on the most recent connection.""" + return set(self.connections[-1].subscriptions) if self.connections else set() + + async def send(self, frames: List[Dict[str, Any]]): + """Send messages to every open, authenticated connection.""" + for c in self.open_connections: + if c.authenticated: + await c.send(frames) + + async def send_status(self, status_: str, message: str): + await self.send([status(status_, message)]) + + async def send_data(self, events: List[Dict[str, Any]]): + """Send data events, e.g. ``[{"ev": "T", "sym": "AAPL", ...}]``.""" + await self.send(events) + + async def close_connections(self, code: int = 1000, reason: str = ""): + """Close every open connection with ``code``.""" + await asyncio.gather(*(c.close(code, reason) for c in self.open_connections)) + + async def drop_connections(self): + """Drop every open connection without a close frame.""" + await asyncio.gather(*(c.drop() for c in self.open_connections)) + + # waiting + + async def wait_for(self, predicate: Callable[[], bool], timeout: float = 5): + """Wait until ``predicate()`` is true, polling every 10 ms.""" + + async def wait(): + while not predicate(): + await asyncio.sleep(0.01) + + await asyncio.wait_for(wait(), timeout) + + async def wait_for_connections(self, n: int, timeout: float = 5): + """Wait until ``n`` clients have authenticated, in total.""" + await self.wait_for( + lambda: sum(c.authenticated for c in self.connections) >= n, timeout + ) + + async def wait_for_subscriptions(self, topics: Set[str], timeout: float = 5): + """Wait until the most recent connection is subscribed to exactly ``topics``.""" + await self.wait_for( + lambda: bool(self.connections) and self.subscriptions == set(topics), + timeout, + ) + + async def wait_for_message( + self, predicate: Callable[[Any], bool], timeout: float = 5 + ) -> Any: + """Wait for a client message matching ``predicate`` and return it.""" + found: List[Any] = [] + + def check() -> bool: + found[:] = [m for m in self.inbox if predicate(m)] + return bool(found) + + await self.wait_for(check, timeout) + return found[0] + + # protocol + + def _process_request( + self, websocket: ServerConnection, request: Request + ) -> Optional[Response]: + if self.reject_handshakes: + return websocket.respond(HTTPStatus.SERVICE_UNAVAILABLE, "unavailable\n") + return None + + async def _handler(self, websocket: ServerConnection): + conn = FakeConnection(self, websocket, len(self.connections)) + self.connections.append(conn) + try: + await conn.send([status("connected", "Connected Successfully")]) + async for raw in websocket: + try: + msg = json.loads(raw) + except ValueError: + msg = raw + conn.inbox.append(msg) + if isinstance(msg, dict): + await self._handle(conn, msg) + if self.on_message is not None: + await self.on_message(conn, msg) + except ConnectionClosed: + pass + finally: + conn.closed.set() + + async def _handle(self, conn: FakeConnection, msg: Dict[str, Any]): + action = msg.get("action") + params = msg.get("params") + if action == "auth": + await self._auth(conn, params) + elif action in ("subscribe", "unsubscribe") and isinstance(params, str): + topics = [p.strip() for p in params.split(",") if p.strip()] + await self._subscription(conn, action, topics) + + async def _auth(self, conn: FakeConnection, key: Optional[str]): + conn.api_key = key + if self._auth_replies: + reply = self._auth_replies.pop(0) + conn.authenticated = any( + f.get("status") == "auth_success" for f in reply.frames + ) + if reply.frames: + await conn.send(reply.frames) + if reply.close_code is not None: + await conn.websocket.close(reply.close_code, reply.close_reason) + return + if key in self.reject_keys or ( + self.valid_keys is not None and key not in self.valid_keys + ): + await conn.send([status("auth_failed", "authentication failed")]) + await conn.websocket.close(1008, "authentication failed") + return + conn.authenticated = True + await conn.send([status("auth_success", "authenticated")]) + + async def _subscription(self, conn: FakeConnection, action: str, topics: List[str]): + replies = [] + for t in topics: + if action == "subscribe" and t in self.reject_topics: + replies.append(status("error", f"subscription rejected: {t}")) + continue + if action == "subscribe": + conn.subscriptions.add(t) + ack = f"subscribed to: {t}" + else: + conn.subscriptions.discard(t) + ack = f"unsubscribed to: {t}" + if self.ack_subscriptions: + replies.append(status("success", ack)) + if replies: + await conn.send(replies) diff --git a/test_websocket/test_fake_server.py b/test_websocket/test_fake_server.py new file mode 100644 index 00000000..9138051e --- /dev/null +++ b/test_websocket/test_fake_server.py @@ -0,0 +1,167 @@ +import asyncio +import threading +import unittest +from concurrent.futures import ThreadPoolExecutor +from typing import Dict, List + +from fake_server import FakeConnection, FakeServer + +from massive import WebSocketClient +from massive.exceptions import AuthError +from massive.websocket import EquityTrade + + +def trade(sym: str, price: float = 1.0): + return {"ev": "T", "sym": sym, "x": 4, "p": price, "s": 100, "t": 1, "q": 1} + + +class FakeServerTest(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self) -> None: + self.server = await FakeServer().start() + self.addAsyncCleanup(self.server.stop) + + def client(self, **kwargs) -> WebSocketClient: + kwargs.setdefault("api_key", "key") + return WebSocketClient(feed=self.server.feed, secure=False, **kwargs) + + async def test_connect_auth_subscribe_and_receive(self) -> None: + c = self.client(subscriptions=["T.AAPL"]) + received: List[EquityTrade] = [] + + async def processor(msgs): + received.extend(msgs) + + task = asyncio.create_task(c.connect(processor)) + await self.server.wait_for_subscriptions({"T.AAPL"}) + conn = self.server.connection + self.assertEqual(conn.path, "/stocks") + self.assertEqual(conn.api_key, "key") + self.assertEqual( + conn.inbox, + [ + {"action": "auth", "params": "key"}, + {"action": "subscribe", "params": "T.AAPL"}, + ], + ) + + await self.server.send_data([trade("AAPL", 2.5)]) + await self.server.wait_for(lambda: len(received) == 1) + self.assertEqual(received[0].symbol, "AAPL") + self.assertEqual(received[0].price, 2.5) + + c.subscribe("T.MSFT") + c.unsubscribe("T.AAPL") + await self.server.wait_for_subscriptions({"T.MSFT"}) + self.assertEqual(c.subs, {"T.MSFT"}) + + await c.close() + await asyncio.wait_for(task, 5) + await self.server.wait_for(lambda: not conn.open) + + async def test_auth_failed_raises(self) -> None: + self.server.reject_keys = {"bad"} + c = self.client(api_key="bad") + with self.assertRaises(AuthError): + await asyncio.wait_for(c.connect(lambda msgs: asyncio.sleep(0)), 5) + self.assertEqual(self.server.inbox, [{"action": "auth", "params": "bad"}]) + + async def test_resubscribes_after_reconnect(self) -> None: + c = self.client(subscriptions=["T.AAPL", "Q.AAPL"], reconnect_backoff=0.01) + task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0))) + await self.server.wait_for_subscriptions({"T.AAPL", "Q.AAPL"}) + + # an abnormal close makes the client reconnect and subscribe again + await self.server.close_connections(code=1011) + await self.server.wait_for_connections(2) + await self.server.wait_for_subscriptions({"T.AAPL", "Q.AAPL"}) + self.assertEqual(len(self.server.open_connections), 1) + + await c.close() + await asyncio.wait_for(task, 5) + + async def test_scripted_auth_reply_and_messages(self) -> None: + # the fake server can replace its replies, e.g. to emulate other statuses + replies: List[str] = [] + + async def on_message(conn: FakeConnection, msg): + if msg.get("action") == "subscribe": + await conn.send([trade(t.split(".", 1)[1]) for t in conn.subscriptions]) + + self.server.on_message = on_message + self.server.ack_subscriptions = False + self.server.script_auth( + {"ev": "status", "status": "auth_success", "message": "scripted"} + ) + c = self.client(subscriptions=["T.AMZN"]) + + async def processor(msgs): + replies.extend(m.symbol for m in msgs) + await c.close() + + await asyncio.wait_for(c.connect(processor), 5) + self.assertEqual(replies, ["AMZN"]) + self.assertTrue(self.server.connection.authenticated) + + async def test_rejected_handshake_is_retried(self) -> None: + self.server.reject_handshakes = True + c = self.client(subscriptions=["T.AAPL"], reconnect_backoff=0.01) + with self.assertLogs("WebSocketClient", level="WARNING") as logs: + task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0))) + # each refused attempt is logged before the next one + await self.server.wait_for(lambda: len(logs.records) >= 2) + self.assertEqual(self.server.connections, []) + self.server.reject_handshakes = False + await self.server.wait_for_subscriptions({"T.AAPL"}) + await c.close() + await asyncio.wait_for(task, 5) + + +class ConcurrentClientsTest(unittest.IsolatedAsyncioTestCase): + """Independent clients, each with its own thread and event loop.""" + + threads = 16 + + async def test_clients_in_parallel_threads(self) -> None: + async def on_subscribe(conn: FakeConnection, msg): + if msg.get("action") == "subscribe": + await conn.send( + [trade(t.split(".", 1)[1]) for t in msg["params"].split(",")] + ) + + server = await FakeServer(on_message=on_subscribe).start() + self.addAsyncCleanup(server.stop) + barrier = threading.Barrier(self.threads, timeout=10) + received: Dict[int, List[str]] = {} + + def run(i: int): + c = WebSocketClient( + api_key=f"key{i}", + feed=server.feed, + secure=False, + subscriptions=[f"T.S{i}"], + ) + + async def processor(msgs): + received[i] = [m.symbol for m in msgs] + await c.close() + + barrier.wait() + asyncio.run(asyncio.wait_for(c.connect(processor), 10)) + + loop = asyncio.get_running_loop() + with ThreadPoolExecutor(self.threads) as pool: + await asyncio.gather( + *(loop.run_in_executor(pool, run, i) for i in range(self.threads)) + ) + + self.assertEqual(received, {i: [f"S{i}"] for i in range(self.threads)}) + self.assertEqual(len(server.connections), self.threads) + by_key = {c.api_key: c for c in server.connections} + for i in range(self.threads): + self.assertEqual( + by_key[f"key{i}"].inbox[1], {"action": "subscribe", "params": f"T.S{i}"} + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test_websocket/test_reconnect.py b/test_websocket/test_reconnect.py new file mode 100644 index 00000000..c6cf46d9 --- /dev/null +++ b/test_websocket/test_reconnect.py @@ -0,0 +1,570 @@ +import asyncio +import json +import socket +import ssl +import threading +import time +import unittest +from typing import Any, Generator, List, Optional +from unittest import mock + +from fake_server import FakeServer, status +from websockets.asyncio.server import ServerConnection +from websockets.datastructures import Headers +from websockets.exceptions import ( + InvalidMessage, + InvalidStatus, + InvalidURI, +) +from websockets.http11 import Request, Response + +import massive.websocket as ws_module +from massive import WebSocketClient +from massive.exceptions import AuthError + +# small waits keep the tests fast; the default is 15 s +BACKOFF = 0.01 +MAX_BACKOFF = 0.05 + + +class TimedServer(FakeServer): + """A FakeServer that records when each handshake was attempted.""" + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.attempts: List[float] = [] + # if set, answer handshakes with this HTTP status + self.reject_status: Optional[int] = None + + def _process_request( + self, websocket: ServerConnection, request: Request + ) -> Optional[Response]: + self.attempts.append(time.monotonic()) + if self.reject_status is not None: + return websocket.respond(self.reject_status, "rejected\n") + return super()._process_request(websocket, request) + + +async def ignore(msgs: Any) -> None: + pass + + +class FakeConnect: + """ + Stands in for websockets' connect(): each attempt waits for `opened` and + then returns `conn`, even if it was cancelled meanwhile, like a handshake + that completes as it is cancelled. + """ + + def __init__(self) -> None: + self.conn = mock.AsyncMock() + self.attempts = 0 + self.started = asyncio.Event() + self.opened = asyncio.Event() + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + fake = self + + class Attempt: + def __await__(self) -> Generator[Any, None, Any]: + return fake._attempt().__await__() + + return Attempt() + + async def _attempt(self) -> Any: + self.attempts += 1 + self.started.set() + try: + await self.opened.wait() + except asyncio.CancelledError: + pass + return self.conn + + +class ReconnectTest(unittest.IsolatedAsyncioTestCase): + """How the client behaves when its connection ends.""" + + async def asyncSetUp(self) -> None: + self.server = TimedServer() + await self.server.start() + self.addAsyncCleanup(self.server.stop) + + def client(self, **kwargs) -> WebSocketClient: + kwargs.setdefault("api_key", "key") + kwargs.setdefault("reconnect_backoff", BACKOFF) + kwargs.setdefault("max_reconnect_backoff", MAX_BACKOFF) + return WebSocketClient(feed=self.server.feed, secure=False, **kwargs) + + def start(self, c: WebSocketClient, processor=ignore) -> "asyncio.Task[None]": + task = asyncio.create_task(c.connect(processor)) + + async def cancel(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + self.addAsyncCleanup(cancel) + return task + + async def test_keeps_changes_scheduled_during_an_outage(self) -> None: + c = self.client(subscriptions=["T.KEEP", "T.OLD", "Q.OLD"]) + task = self.start(c) + await self.server.wait_for_subscriptions({"T.KEEP", "T.OLD", "Q.OLD"}) + + self.server.reject_handshakes = True + # scheduled but not sent yet when the connection drops + c.subscribe("T.NEW") + c.unsubscribe("T.OLD") + await self.server.drop_connections() + # scheduled while the client cannot reconnect + c.subscribe("Q.NEW") + c.unsubscribe("Q.OLD") + self.server.reject_handshakes = False + + await self.server.wait_for_connections(2) + await self.server.wait_for_subscriptions({"T.KEEP", "T.NEW", "Q.NEW"}) + topics: List[str] = [] + for m in self.server.connection.inbox[1:]: + self.assertEqual(m["action"], "subscribe") + topics += m["params"].split(",") + self.assertCountEqual(topics, ["T.KEEP", "T.NEW", "Q.NEW"]) + self.assertEqual(c.scheduled_subs, {"T.KEEP", "T.NEW", "Q.NEW"}) + + await c.close() + await asyncio.wait_for(task, 5) + + async def test_reconnects_after_a_clean_close_until_closed(self) -> None: + c = self.client(subscriptions=["T.AAPL"]) + task = self.start(c) + await self.server.wait_for_subscriptions({"T.AAPL"}) + + # e.g. the server closes cleanly after a slow_consumer status + await self.server.close_connections(code=1000) + await self.server.wait_for_connections(2) + await self.server.wait_for_subscriptions({"T.AAPL"}) + self.assertFalse(task.done()) + + await c.close() + await asyncio.wait_for(task, 5) + self.assertEqual(len(self.server.connections), 2) + + async def test_waits_between_attempts(self) -> None: + c = self.client( + reconnect_backoff=0.05, max_reconnect_backoff=0.2, max_reconnects=None + ) + delays: List[float] = [] + delay = c._reconnect_delay + + def record(attempt: int) -> float: + delays.append(delay(attempt)) + return delays[-1] + + with mock.patch.object(c, "_reconnect_delay", side_effect=record): + task = self.start(c) + await self.server.wait_for_connections(1) + self.server.reject_handshakes = True + await self.server.drop_connections() + # 0.05, 0.1, 0.2, then the cap + await self.server.wait_for(lambda: len(self.server.attempts) == 6) + + # close() during the wait stops connect() + await c.close() + await asyncio.wait_for(task, 5) + + self.assertEqual(len(self.server.attempts), 6) + # one delay before each refused attempt, maybe one more being waited + self.assertIn(len(delays), (5, 6)) + for d in delays: + self.assertGreaterEqual(d, 0.05) + self.assertLessEqual(d, 0.2) + self.assertGreaterEqual(delays[1], 0.1) + # at the cap, jitter cannot lengthen the wait + self.assertEqual(delays[2:], [0.2] * (len(delays) - 2)) + # each refused attempt waited for its delay before the next one + gaps = [ + b - a for a, b in zip(self.server.attempts[1:], self.server.attempts[2:]) + ] + # allow for the timer resolution, about 15.6 ms on Windows + for gap, d in zip(gaps, delays[1:]): + self.assertGreaterEqual(gap, d - 0.03) + + async def test_transient_handshake_errors_are_retried(self) -> None: + for code in (408, 429, 500, 503): + self.server.reject_status = code + n = len(self.server.attempts) + c = self.client() + with self.assertLogs("WebSocketClient", level="WARNING") as logs: + task = self.start(c) + await self.server.wait_for(lambda: len(self.server.attempts) >= n + 3) + self.assertIn( + f"connect failed; reconnecting in 0.0 s: server rejected WebSocket " + f"connection: HTTP {code}", + logs.output[0], + ) + self.server.reject_status = None + await self.server.wait_for( + lambda: bool(self.server.open_connections), timeout=5 + ) + await c.close() + await asyncio.wait_for(task, 5) + + async def test_refused_connections_are_retried(self) -> None: + port = self.server.port + await self.server.stop() + c = WebSocketClient( + api_key="key", feed=f"127.0.0.1:{port}", secure=False, reconnect_backoff=0 + ) + with self.assertLogs("WebSocketClient", level="WARNING") as logs: + task = self.start(c) + # Windows retries the SYN before it reports a refusal, so there an + # attempt can take about 2 s + await self.server.wait_for(lambda: len(logs.records) >= 2, timeout=15) + self.assertFalse(task.done()) + self.assertIn("connect failed; reconnecting in", logs.output[0]) + + async def test_fatal_handshake_errors_raise_at_once(self) -> None: + for code in (400, 401, 403, 404): + self.server.reject_status = code + n = len(self.server.attempts) + with self.assertRaises(InvalidStatus) as e: + await asyncio.wait_for(self.client().connect(ignore), 5) + self.assertEqual(e.exception.response.status_code, code) + self.assertEqual(len(self.server.attempts), n + 1) + + async def test_invalid_arguments_raise_at_once(self) -> None: + # rejected when building the connection + with self.assertRaises(ValueError): + await self.client(max_reconnects=None).connect(ignore, compression="x") + # rejected by the first attempt + with self.assertRaises(TypeError): + await self.client(max_reconnects=None).connect(ignore, bogus=1) + # connect() decides which errors to retry + with self.assertRaises(TypeError): + await self.client().connect(ignore, process_exception=lambda e: None) + c = WebSocketClient(api_key="key", feed="127.0.0.1:port", secure=False) + with self.assertRaises(ValueError): + await asyncio.wait_for(c.connect(ignore), 5) + self.assertEqual(self.server.attempts, []) + + async def test_outages_with_successes_between_do_not_raise(self) -> None: + c = self.client(subscriptions=["T.AAPL"], max_reconnects=5) + task = self.start(c) + for i in range(1, 8): + await self.server.wait_for_connections(i) + await self.server.wait_for_subscriptions({"T.AAPL"}) + await self.server.drop_connections() + await self.server.wait_for_connections(8) + self.assertFalse(task.done()) + + await c.close() + await asyncio.wait_for(task, 5) + + async def test_consecutive_failures_raise(self) -> None: + c = self.client(max_reconnects=2) + task = self.start(c) + await self.server.wait_for_connections(1) + self.server.reject_handshakes = True + await self.server.drop_connections() + + # the drop and two refused attempts + with self.assertRaises(InvalidStatus): + await asyncio.wait_for(task, 5) + self.assertEqual(len(self.server.attempts), 3) + + async def test_unlimited_reconnects(self) -> None: + c = self.client(max_reconnects=None, max_reconnect_backoff=BACKOFF) + task = self.start(c) + self.server.reject_handshakes = True + await self.server.wait_for(lambda: len(self.server.attempts) >= 20) + self.assertFalse(task.done()) + self.server.reject_handshakes = False + await self.server.wait_for_connections(1) + + await c.close() + await asyncio.wait_for(task, 5) + + async def request_close_from_a_thread(self, c: WebSocketClient): + thread = threading.Thread(target=c.request_close) + thread.start() + await asyncio.get_running_loop().run_in_executor(None, thread.join) + + async def test_request_close_from_another_thread(self) -> None: + c = self.client(subscriptions=["T.AAPL"]) + task = self.start(c) + await self.server.wait_for_subscriptions({"T.AAPL"}) + # reconnects after a clean close, until request_close() + await self.server.close_connections(code=1000) + await self.server.wait_for_connections(2) + await self.server.wait_for_subscriptions({"T.AAPL"}) + + await self.request_close_from_a_thread(c) + await asyncio.wait_for(task, 2) + await self.server.wait_for(lambda: not self.server.open_connections) + self.assertEqual(len(self.server.connections), 2) + + async def test_request_close_while_waiting_to_reconnect(self) -> None: + c = self.client(reconnect_backoff=60, max_reconnect_backoff=60) + task = self.start(c) + await self.server.wait_for_connections(1) + await self.server.drop_connections() + + await self.request_close_from_a_thread(c) + await asyncio.wait_for(task, 2) + self.assertEqual(len(self.server.attempts), 1) + + async def test_request_close_while_connecting(self) -> None: + # a server that accepts TCP connections but never answers the handshake + accepted = asyncio.Event() + + async def stall(reader: asyncio.StreamReader, writer: asyncio.StreamWriter): + accepted.set() + await reader.read() + writer.close() + + server = await asyncio.start_server(stall, "127.0.0.1", 0) + self.addAsyncCleanup(server.wait_closed) + self.addCleanup(server.close) + port = server.sockets[0].getsockname()[1] + c = WebSocketClient(api_key="key", feed=f"127.0.0.1:{port}", secure=False) + task = self.start(c) + await asyncio.wait_for(accepted.wait(), 5) + + await self.request_close_from_a_thread(c) + await asyncio.wait_for(task, 2) + + async def test_cancel_while_connecting_closes_the_socket(self) -> None: + accepted = asyncio.Event() + disconnected = asyncio.Event() + + async def stall(reader: asyncio.StreamReader, writer: asyncio.StreamWriter): + accepted.set() + await reader.read() + disconnected.set() + writer.close() + + server = await asyncio.start_server(stall, "127.0.0.1", 0) + self.addAsyncCleanup(server.wait_closed) + self.addCleanup(server.close) + port = server.sockets[0].getsockname()[1] + c = WebSocketClient(api_key="key", feed=f"127.0.0.1:{port}", secure=False) + task = asyncio.create_task(c.connect(ignore)) + await asyncio.wait_for(accepted.wait(), 5) + + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + await asyncio.wait_for(disconnected.wait(), 2) + + async def test_cancel_closes_a_connection_opened_meanwhile(self) -> None: + fake = FakeConnect() + with mock.patch.object(ws_module, "connect", fake): + task = asyncio.create_task(self.client().connect(ignore)) + await asyncio.wait_for(fake.started.wait(), 5) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + self.assertEqual(fake.attempts, 1) + fake.conn.close.assert_awaited_once() + + async def test_no_attempt_after_request_close(self) -> None: + fake = FakeConnect() + with mock.patch.object(ws_module, "connect", fake): + c = self.client() + task = asyncio.create_task(c.connect(ignore)) + await asyncio.wait_for(fake.started.wait(), 5) + c.request_close() + # the handshake completes as the close is requested + fake.opened.set() + await asyncio.wait_for(task, 5) + # give a stray attempt every chance to start + for _ in range(20): + await asyncio.sleep(0) + self.assertEqual(fake.attempts, 1) + fake.conn.close.assert_awaited_once() + fake.conn.recv.assert_not_awaited() + + async def test_request_close_before_connect(self) -> None: + c = self.client() + c.request_close() + await asyncio.wait_for(c.connect(ignore), 2) + self.assertEqual(self.server.attempts, []) + + # the request applied to that connect() only + task = self.start(c) + await self.server.wait_for_connections(1) + c.request_close() + await asyncio.wait_for(task, 2) + + async def test_close_request_stays_pending_if_connect_raises(self) -> None: + c = self.client() + c.request_close() + with self.assertRaises(ValueError): + await c.connect(ignore, compression="x") + # the next connect() returns because of it, which consumes it + await asyncio.wait_for(c.connect(ignore), 2) + self.assertEqual(self.server.attempts, []) + task = self.start(c) + await self.server.wait_for_connections(1) + c.request_close() + await asyncio.wait_for(task, 2) + + +class TransientTest(unittest.TestCase): + """Which failed attempts to connect are retried.""" + + def test_classification(self) -> None: + def status(code: int) -> InvalidStatus: + return InvalidStatus(Response(code, "", Headers(), b"")) + + def invalid(cause: Optional[BaseException]) -> InvalidMessage: + # how websockets 15 reports a handshake without a valid response + e = InvalidMessage("did not receive a valid HTTP response") + e.__cause__ = cause + return e + + transient = [ + status(408), + status(429), + status(500), + status(503), + invalid(EOFError("connection closed while reading HTTP status line")), + invalid(ConnectionResetError()), + # how websockets 14 reports a connection closed before a response + EOFError("line without CRLF"), + ConnectionRefusedError(), + OSError("unreachable"), + socket.gaierror(), + TimeoutError(), + asyncio.TimeoutError(), + ] + fatal = [ + status(400), + status(401), + status(403), + status(404), + invalid(ValueError("invalid HTTP status line")), + invalid(None), + InvalidURI("ws://x", "bad"), + ssl.SSLCertVerificationError(), + TypeError(), + ValueError(), + AuthError("bad key"), + ] + for e in transient: + self.assertTrue(WebSocketClient._is_transient(e), repr(e)) + for e in fatal: + self.assertFalse(WebSocketClient._is_transient(e), repr(e)) + + +class MaxConnectionsTest(unittest.IsolatedAsyncioTestCase): + """A max_connections status in the reply to auth reaches the caller.""" + + async def asyncSetUp(self) -> None: + self.server = await FakeServer().start() + self.addAsyncCleanup(self.server.stop) + self.status = status( + "max_connections", "Maximum number of websocket connections exceeded." + ) + + async def test_raw_processor_receives_it(self) -> None: + self.server.script_auth(self.status, close_code=1008) + c = WebSocketClient( + api_key="key", + feed=self.server.feed, + secure=False, + raw=True, + reconnect_backoff=BACKOFF, + subscriptions=["T.AAPL"], + ) + frames: List[Any] = [] + + async def processor(msg: Any) -> None: + frames.append((json.loads(msg), c.max_connections_reached)) + + task = asyncio.create_task(c.connect(processor)) + # it reconnects after the server closes, and the next handshake succeeds + await self.server.wait_for_connections(1) + await self.server.wait_for_subscriptions({"T.AAPL"}) + self.assertEqual(len(self.server.connections), 2) + self.assertEqual(frames[0], ([self.status], True)) + self.assertFalse(c.max_connections_reached) + + c.request_close() + await asyncio.wait_for(task, 5) + + async def test_flag_is_set(self) -> None: + self.server.script_auth(self.status, close_code=1008) + self.server.script_auth(self.status, close_code=1008) + c = WebSocketClient( + api_key="key", + feed=self.server.feed, + secure=False, + reconnect_backoff=0.2, + max_reconnect_backoff=0.2, + ) + delays: List[int] = [] + delay = c._reconnect_delay + + def record(attempt: int) -> float: + delays.append(attempt) + return delay(attempt) + + with mock.patch.object(c, "_reconnect_delay", side_effect=record): + task = asyncio.create_task(c.connect(ignore)) + await self.server.wait_for(lambda: c.max_connections_reached) + # cleared by the next successful handshake + await self.server.wait_for_connections(1) + await self.server.wait_for(lambda: not c.max_connections_reached) + self.assertEqual(len(self.server.connections), 3) + c.request_close() + await asyncio.wait_for(task, 5) + # a refused handshake is a failed attempt, so the backoff grows + self.assertEqual(delays, [0, 1]) + + +class ReconnectDelayTest(unittest.TestCase): + """The wait before each consecutive reconnect attempt.""" + + def delays(self, c: WebSocketClient, n: int) -> List[float]: + return [c._reconnect_delay(i) for i in range(n)] + + def test_defaults(self) -> None: + c = WebSocketClient(api_key="key") + self.assertEqual(c.reconnect_backoff, 15.0) + self.assertEqual(c.max_reconnect_backoff, 120.0) + self.assertIsNone(c.max_reconnects) + with mock.patch.object(ws_module.random, "uniform", side_effect=lambda a, b: a): + self.assertEqual(self.delays(c, 6), [15, 30, 60, 120, 120, 120]) + with mock.patch.object(ws_module.random, "uniform", side_effect=lambda a, b: b): + self.assertEqual(self.delays(c, 6), [18, 36, 72, 120, 120, 120]) + + def test_jitter_is_clamped_to_the_cap(self) -> None: + c = WebSocketClient( + api_key="key", reconnect_backoff=10, max_reconnect_backoff=100 + ) + with mock.patch.object(ws_module.random, "uniform", side_effect=lambda a, b: a): + self.assertEqual(self.delays(c, 6), [10, 20, 40, 80, 100, 100]) + with mock.patch.object(ws_module.random, "uniform", side_effect=lambda a, b: b): + self.assertEqual(self.delays(c, 6), [12, 24, 48, 96, 100, 100]) + + def test_every_delay_is_within_bounds(self) -> None: + c = WebSocketClient(api_key="key") + for attempt in list(range(20)) * 50 + [10**6]: + d = c._reconnect_delay(attempt) + self.assertGreaterEqual(d, 15.0) + self.assertLessEqual(d, 120.0) + + def test_equal_bounds(self) -> None: + c = WebSocketClient(api_key="key", reconnect_backoff=0, max_reconnect_backoff=0) + self.assertEqual(self.delays(c, 3), [0, 0, 0]) + + def test_invalid_bounds(self) -> None: + with self.assertRaises(ValueError): + WebSocketClient(api_key="key", reconnect_backoff=-1) + with self.assertRaises(ValueError): + WebSocketClient( + api_key="key", reconnect_backoff=30, max_reconnect_backoff=20 + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test_websocket/test_status.py b/test_websocket/test_status.py index d333c9c8..1212e38e 100644 --- a/test_websocket/test_status.py +++ b/test_websocket/test_status.py @@ -3,6 +3,7 @@ from base_ws import BaseTest from mock_server import port from websockets.asyncio.client import connect +from websockets.exceptions import ConnectionClosedOK import asyncio @@ -65,11 +66,15 @@ async def test_success_statuses_are_not_logged(self): self.assertEqual(logs.output, ["WARNING:WebSocketClient:sentinel"]) async def test_disconnect_reason_is_logged(self): - # the server sends a status and then closes the connection normally + # the server sends a status and then closes the connection normally; + # the client would reconnect, so let it give up instead with self.assertLogs("WebSocketClient", level="WARNING") as logs: - c = WebSocketClient(api_key="", feed=f"localhost:{port}", secure=False) + c = WebSocketClient( + api_key="", feed=f"localhost:{port}", secure=False, max_reconnects=0 + ) c.subscribe("SLOW.X") - await asyncio.wait_for(c.connect(lambda msgs: asyncio.sleep(0)), 5) + with self.assertRaises(ConnectionClosedOK): + await asyncio.wait_for(c.connect(lambda msgs: asyncio.sleep(0)), 5) self.assertEqual( logs.output, [