From ba72dfbe3a4ab7a5c8c91563c88dd33fcac8a738 Mon Sep 17 00:00:00 2001 From: himanshu Date: Tue, 27 Oct 2026 16:02:24 +0530 Subject: [PATCH] Handle WebSocket error frames and fix reconnect retries; bump version to 0.3.5 --- src/newsdataapi/__init__.py | 2 +- src/newsdataapi/websocket.py | 69 +++++++++++++---- tests/test_websocket.py | 145 +++++++++++++++++++++++++++++++++-- 3 files changed, 191 insertions(+), 25 deletions(-) diff --git a/src/newsdataapi/__init__.py b/src/newsdataapi/__init__.py index a01c6f4..31b291c 100644 --- a/src/newsdataapi/__init__.py +++ b/src/newsdataapi/__init__.py @@ -17,7 +17,7 @@ ) from .websocket import NewsDataApiWebSocket -__version__ = "0.3.4" +__version__ = "0.3.5" __all__ = [ "NewsDataApiClient", diff --git a/src/newsdataapi/websocket.py b/src/newsdataapi/websocket.py index 148bf0f..4e868b5 100644 --- a/src/newsdataapi/websocket.py +++ b/src/newsdataapi/websocket.py @@ -30,7 +30,7 @@ if _WEBSOCKET_SUPPORTED: from websockets.asyncio import client as _asyncio_client - from websockets.exceptions import ConnectionClosedError, InvalidStatus + from websockets.exceptions import ConnectionClosedError, InvalidMessage, InvalidStatus from websockets.sync import client as _sync_client logger = logging.getLogger(__name__) @@ -44,13 +44,16 @@ def _check_registration_id(registration_id: str) -> None: ) -def _permanent_auth_error(exc: Exception) -> NewsdataWebSocketAuthError | None: +def _permanent_auth_error( + exc: Exception, error_response: dict[str, Any] | None +) -> NewsdataWebSocketAuthError | None: """Return the auth error to raise if ``exc`` is a permanent rejection. Returns ``None`` for transient failures (which the caller reconnects on): handshake 401/403 and policy-violation close 1008 are permanent; anything else (other handshake status, other close codes, network ``OSError``) is - transient. + transient. ``error_response`` is the error response the server sent just + before closing, if any; it becomes the exception's argument. """ if isinstance(exc, InvalidStatus): if exc.response.status_code in (401, 403): @@ -59,14 +62,20 @@ def _permanent_auth_error(exc: Exception) -> NewsdataWebSocketAuthError | None: if isinstance(exc, ConnectionClosedError): close = exc.rcvd if close is not None and close.code == constants.WS_POLICY_VIOLATION: + if error_response is not None: + return NewsdataWebSocketAuthError(error_response) return NewsdataWebSocketAuthError(close.reason or "connection rejected") return None return None -def _transient_error(exc: Exception) -> NewsdataWebSocketError: +def _transient_error( + exc: Exception, error_response: dict[str, Any] | None +) -> NewsdataWebSocketError: """Wrap a transient failure as a :class:`NewsdataWebSocketError` (used only when ``reconnect=False`` so the caller stops instead of retrying).""" + if error_response is not None: + return NewsdataWebSocketError(error_response) if isinstance(exc, InvalidStatus): return NewsdataWebSocketError(f"handshake failed (HTTP {exc.response.status_code})") if isinstance(exc, ConnectionClosedError): @@ -295,24 +304,38 @@ def stream(self, registration_id: str) -> Iterator[dict[str, Any]]: url = self._url(registration_id) delay = self._reconnect_delay while True: + error_response: dict[str, Any] | None = None try: with _sync_client.connect(url, **self._connect_kwargs()) as websocket: self._ws_sync = websocket logger.info("connected to %s", _redact_url(url)) - delay = self._reconnect_delay # reset after a successful connect + # reset after a successful connect, unless the server refuses it + refused_delay, delay = delay, self._reconnect_delay for message in websocket: try: response = json.loads(message) except ValueError: continue # skip malformed frames - if isinstance(response, dict): - yield response - except (InvalidStatus, ConnectionClosedError, OSError) as exc: - auth = _permanent_auth_error(exc) + if not isinstance(response, dict): + continue + if response.get("status") == "error": + # sent right before the server closes the connection: + # raise it with that close, and keep backing off + error_response, delay = response, refused_delay + continue + yield response + except ( + InvalidStatus, + InvalidMessage, # connection lost during the handshake + ConnectionClosedError, + OSError, + asyncio.TimeoutError, # not an OSError on Python 3.10 + ) as exc: + auth = _permanent_auth_error(exc, error_response) if auth is not None: raise auth from exc if not self._reconnect: - raise _transient_error(exc) from exc + raise _transient_error(exc, error_response) from exc logger.warning( "connection failed (%s); reconnecting in %.1fs", exc, delay ) @@ -330,26 +353,40 @@ async def stream_async(self, registration_id: str) -> AsyncIterator[dict[str, An url = self._url(registration_id) delay = self._reconnect_delay while True: + error_response: dict[str, Any] | None = None try: async with _asyncio_client.connect( url, **self._connect_kwargs() ) as websocket: self._ws_async = websocket logger.info("connected to %s", _redact_url(url)) - delay = self._reconnect_delay # reset after a successful connect + # reset after a successful connect, unless the server refuses it + refused_delay, delay = delay, self._reconnect_delay async for message in websocket: try: response = json.loads(message) except ValueError: continue # skip malformed frames - if isinstance(response, dict): - yield response - except (InvalidStatus, ConnectionClosedError, OSError) as exc: - auth = _permanent_auth_error(exc) + if not isinstance(response, dict): + continue + if response.get("status") == "error": + # sent right before the server closes the connection: + # raise it with that close, and keep backing off + error_response, delay = response, refused_delay + continue + yield response + except ( + InvalidStatus, + InvalidMessage, # connection lost during the handshake + ConnectionClosedError, + OSError, + asyncio.TimeoutError, # not an OSError on Python 3.10 + ) as exc: + auth = _permanent_auth_error(exc, error_response) if auth is not None: raise auth from exc if not self._reconnect: - raise _transient_error(exc) from exc + raise _transient_error(exc, error_response) from exc logger.warning( "connection failed (%s); reconnecting in %.1fs", exc, delay ) diff --git a/tests/test_websocket.py b/tests/test_websocket.py index b42b33d..009fd19 100644 --- a/tests/test_websocket.py +++ b/tests/test_websocket.py @@ -12,7 +12,7 @@ import pytest from websockets.datastructures import Headers -from websockets.exceptions import ConnectionClosedError, InvalidStatus +from websockets.exceptions import ConnectionClosedError, InvalidMessage, InvalidStatus from websockets.frames import Close from websockets.http11 import Response @@ -39,6 +39,11 @@ def _news(*titles: str) -> str: ) +def _error(code: str, message: str = "rejected") -> str: + """An error frame, shaped like the server's error response.""" + return json.dumps({"status": "error", "results": {"message": message, "code": code}}) + + def _closed(code: int, reason: str = "") -> ConnectionClosedError: return ConnectionClosedError(Close(code, reason), None) @@ -47,6 +52,11 @@ def _invalid_status(status: int) -> InvalidStatus: return InvalidStatus(Response(status, "x", Headers(), b"")) +# Terminals the fakes raise from connect() itself (handshake / network +# failures) rather than during iteration. +_CONNECT_ERRORS = (InvalidStatus, InvalidMessage, OSError, asyncio.TimeoutError) + + class _FakeWS: """A scripted sync connection: yields ``messages`` then raises ``terminal`` (a ConnectionClosedError) or ends normally if ``terminal`` is None.""" @@ -75,9 +85,10 @@ def close(self) -> None: class _FakeConnect: """Scripted ``connect``: each attempt is ``(messages, terminal)``. - ``terminal`` raised from ``connect()`` itself if it is an ``InvalidStatus`` - or ``OSError`` (handshake / network failure), otherwise during iteration. - Once the script is exhausted it closes with 1008 to stop any reconnect loop. + ``terminal`` raised from ``connect()`` itself if it is one of + ``_CONNECT_ERRORS`` (handshake / network failure), otherwise during + iteration. Once the script is exhausted it closes with 1008 to stop any + reconnect loop. """ def __init__(self, *attempts: tuple[list, BaseException | None]) -> None: @@ -95,7 +106,7 @@ def __call__(self, url: str, **kwargs: object) -> _FakeWS: messages, terminal = self._attempts.pop(0) else: messages, terminal = [], _closed(1008, "exhausted") - if isinstance(terminal, (InvalidStatus, OSError)): + if isinstance(terminal, _CONNECT_ERRORS): raise terminal conn = _FakeWS(messages, terminal) self.conns.append(conn) @@ -239,6 +250,18 @@ def test_close_1008_raises_auth_with_reason(monkeypatch: pytest.MonkeyPatch) -> assert fake.calls == 1 +def test_error_response_raises_auth_with_it(monkeypatch: pytest.MonkeyPatch) -> None: + error = _error("ApiLimitExceeded", "out of credits") + fake = _FakeConnect(([_news("A"), error], _closed(1008, "ApiLimitExceeded"))) + monkeypatch.setattr(CONNECT, fake) + got: list[str] = [] + with pytest.raises(NewsdataWebSocketAuthError) as exc_info: + _collect(_ws(), got) # reconnect=True default + assert got == ["A"] # the error response is raised, never yielded + assert exc_info.value.args[0] == json.loads(error) + assert fake.calls == 1 + + # --- reconnect=False: stop on first drop ------------------------------------ @@ -269,6 +292,17 @@ def test_no_reconnect_handshake_5xx_raises(monkeypatch: pytest.MonkeyPatch) -> N assert not isinstance(exc_info.value, NewsdataWebSocketAuthError) +def test_no_reconnect_server_error_response_raises(monkeypatch: pytest.MonkeyPatch) -> None: + error = _error("ServerError", "try again later") + fake = _FakeConnect(([error], _closed(1011, "ServerError"))) + monkeypatch.setattr(CONNECT, fake) + with pytest.raises(NewsdataWebSocketError) as exc_info: + list(_ws(reconnect=False).stream("r")) + assert not isinstance(exc_info.value, NewsdataWebSocketAuthError) + assert exc_info.value.args[0] == json.loads(error) + assert fake.calls == 1 + + # --- reconnect=True: retry transient drops ---------------------------------- @@ -324,6 +358,40 @@ def test_reconnects_on_network_error(monkeypatch: pytest.MonkeyPatch) -> None: assert fake.calls == 2 +def test_reconnects_on_handshake_connection_lost(monkeypatch: pytest.MonkeyPatch) -> None: + fake = _FakeConnect( + ([], InvalidMessage("did not receive a valid HTTP response")), # lost mid-handshake + ([_news("A")], _closed(1008, "stop")), + ) + monkeypatch.setattr(CONNECT, fake) + got: list[str] = [] + with pytest.raises(NewsdataWebSocketAuthError): + _collect(_ws(), got) + assert got == ["A"] + assert fake.calls == 2 + + +def test_server_error_responses_keep_backing_off(monkeypatch: pytest.MonkeyPatch) -> None: + slept: list[float] = [] + monkeypatch.setattr( + "newsdataapi.websocket.time.sleep", lambda d: slept.append(d) + ) + refused = ([_error("ServerError")], _closed(1011, "ServerError")) + fake = _FakeConnect( + refused, + refused, + refused, + ([_news("A")], _closed(1013, "SendTimeout")), # accepted -> backoff starts over + ([_news("B")], _closed(1008, "stop")), + ) + monkeypatch.setattr(CONNECT, fake) + got: list[str] = [] + with pytest.raises(NewsdataWebSocketAuthError): + _collect(_ws(), got) + assert got == ["A", "B"] # error responses are never yielded + assert slept == [1.0, 2.0, 4.0, 1.0] + + # =========================================================================== # Async streaming (NewsDataApiWebSocket.stream_async) # =========================================================================== @@ -373,8 +441,8 @@ async def __aexit__(self, *exc: object) -> bool: class _FakeAsyncConnect: """Scripted async ``connect``: each attempt is ``(messages, terminal)``. - An ``InvalidStatus`` / ``OSError`` terminal is raised from ``__aenter__`` - (handshake failure); otherwise it is raised during async iteration. + A ``_CONNECT_ERRORS`` terminal is raised from ``__aenter__`` (handshake + failure); otherwise it is raised during async iteration. """ def __init__(self, *attempts: tuple[list, BaseException | None]) -> None: @@ -392,7 +460,7 @@ def __call__(self, url: str, **kwargs: object) -> _FakeAsyncCM: messages, terminal = self._attempts.pop(0) else: messages, terminal = [], _closed(1008, "exhausted") - if isinstance(terminal, (InvalidStatus, OSError)): + if isinstance(terminal, _CONNECT_ERRORS): return _FakeAsyncCM(None, terminal) conn = _FakeAsyncWS(messages, terminal) self.conns.append(conn) @@ -502,6 +570,18 @@ def test_async_close_1008_raises_auth_with_reason(monkeypatch: pytest.MonkeyPatc assert fake.calls == 1 +def test_async_error_response_raises_auth_with_it(monkeypatch: pytest.MonkeyPatch) -> None: + error = _error("ApiLimitExceeded", "out of credits") + fake = _FakeAsyncConnect(([_news("A"), error], _closed(1008, "ApiLimitExceeded"))) + monkeypatch.setattr(CONNECT_ASYNC, fake) + got: list[str] = [] + with pytest.raises(NewsdataWebSocketAuthError) as exc_info: + _acollect(_ws(), got) # reconnect=True default + assert got == ["A"] # the error response is raised, never yielded + assert exc_info.value.args[0] == json.loads(error) + assert fake.calls == 1 + + # --- reconnect=False -------------------------------------------------------- @@ -542,3 +622,52 @@ def test_async_reconnects_on_network_error(monkeypatch: pytest.MonkeyPatch) -> N _acollect(_ws(), got) assert got == ["A"] assert fake.calls == 2 + + +def test_async_reconnects_on_handshake_connection_lost(monkeypatch: pytest.MonkeyPatch) -> None: + fake = _FakeAsyncConnect( + ([], InvalidMessage("did not receive a valid HTTP response")), # lost mid-handshake + ([_news("A")], _closed(1008, "stop")), + ) + monkeypatch.setattr(CONNECT_ASYNC, fake) + got: list[str] = [] + with pytest.raises(NewsdataWebSocketAuthError): + _acollect(_ws(), got) + assert got == ["A"] + assert fake.calls == 2 + + +def test_async_reconnects_on_open_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + fake = _FakeAsyncConnect( + ([], asyncio.TimeoutError()), # handshake timeout; not an OSError on Python 3.10 + ([_news("A")], _closed(1008, "stop")), + ) + monkeypatch.setattr(CONNECT_ASYNC, fake) + got: list[str] = [] + with pytest.raises(NewsdataWebSocketAuthError): + _acollect(_ws(), got) + assert got == ["A"] + assert fake.calls == 2 + + +def test_async_server_error_responses_keep_backing_off(monkeypatch: pytest.MonkeyPatch) -> None: + slept: list[float] = [] + + async def record(delay: float) -> None: + slept.append(delay) + + monkeypatch.setattr("newsdataapi.websocket.asyncio.sleep", record) + refused = ([_error("ServerError")], _closed(1011, "ServerError")) + fake = _FakeAsyncConnect( + refused, + refused, + refused, + ([_news("A")], _closed(1013, "SendTimeout")), # accepted -> backoff starts over + ([_news("B")], _closed(1008, "stop")), + ) + monkeypatch.setattr(CONNECT_ASYNC, fake) + got: list[str] = [] + with pytest.raises(NewsdataWebSocketAuthError): + _acollect(_ws(), got) + assert got == ["A", "B"] # error responses are never yielded + assert slept == [1.0, 2.0, 4.0, 1.0]