Handle WebSocket error frames and fix reconnect retries; bump version…

src/newsdataapi/__init__.py

@@ -17,7 +17,7 @@
 )
 from .websocket import NewsDataApiWebSocket
 
-__version__ = "0.3.4"
+__version__ = "0.3.5"
 
 __all__ = [
     "NewsDataApiClient",

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
                 )

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]