Pass subscription errors to an on_error callback

src/rivian/rivian.py

@@ -1012,6 +1012,7 @@ async def _start_subscription(
         payload: dict[str, Any],
         callback: Callable[[dict[str, Any]], None],
         description: str,
+        on_error: Callable[[dict[str, Any]], None] | None = None,
     ) -> Callable[[], Awaitable[None]] | None:
         """Connect the web socket if needed and start a subscription.
 
@@ -1022,7 +1023,9 @@ async def _start_subscription(
             ws_monitor = await self._ws_connect()
             # Ends early if the connection closes, e.g. when rate limited
             if await ws_monitor.wait_for_ack() and (
-                unsubscribe := await ws_monitor.start_subscription(payload, callback)
+                unsubscribe := await ws_monitor.start_subscription(
+                    payload, callback, on_error=on_error
+                )
             ):
                 _LOGGER.debug("Subscribed to %s", description)
                 return unsubscribe

src/rivian/ws_monitor.py

@@ -98,8 +98,14 @@ def __init__(
         self._monitor_task: asyncio.Task | None = None
         self._receiver_task: asyncio.Task | None = None
         self._last_received: datetime | None = None
+        # id -> (callback, payload, error callback)
         self._subscriptions: dict[
-            str, tuple[Callable[[dict[str, Any]], None], dict[str, Any]]
+            str,
+            tuple[
+                Callable[[dict[str, Any]], None],
+                dict[str, Any],
+                Callable[[dict[str, Any]], None] | None,
+            ],
         ] = {}
 
     @property
@@ -211,13 +217,20 @@ async def wait_for_ack(self) -> bool:
         return self._connection_ack.is_set()
 
     async def start_subscription(
-        self, payload: dict[str, Any], callback: Callable[[dict[str, Any]], None]
+        self,
+        payload: dict[str, Any],
+        callback: Callable[[dict[str, Any]], None],
+        on_error: Callable[[dict[str, Any]], None] | None = None,
     ) -> Callable[[], Awaitable[None]] | None:
-        """Start a subscription."""
+        """Start a subscription.
+
+        `on_error` gets the subscription's error messages, which are otherwise
+        only logged.
+        """
         if not self.connected:
             return None
         _id = str(uuid4())
-        self._subscriptions[_id] = (callback, payload)
+        self._subscriptions[_id] = (callback, payload, on_error)
         await self._subscribe(_id, payload)
 
         async def unsubscribe() -> None:
@@ -247,7 +260,7 @@ async def _resubscribe_all(self) -> None:
         except TimeoutError:
             _LOGGER.error("A timeout occurred while attempting to resubscribe")
             return
-        for _id, (_, payload) in self._subscriptions.items():
+        for _id, (_, payload, _) in self._subscriptions.items():
             await self._subscribe(_id, payload)
 
     async def _receiver(self) -> None:
@@ -305,6 +318,12 @@ async def _receive(self, websocket: ClientWebSocketResponse[bool]) -> None:
                                     _fn(data)
                             except Exception:  # pylint: disable=broad-except
                                 _LOGGER.exception("Error in a subscription callback")
+                    elif (
+                        data_type == "error"
+                        and (subscription := self._subscriptions.get(data.get("id")))
+                        and subscription[2]
+                    ):
+                        subscription[2](data)
                     else:
                         self._log_message(msg)
                 elif msg.type == WSMsgType.ERROR:

tests/test_ws_monitor.py

@@ -31,6 +31,7 @@ def __init__(self) -> None:
         self.ping_once_connections = 0
         self.connections = 0
         self.next_on_subscribe = False
+        self.error_on_subscribe = False
         self.url = ""
 
     async def handler(self, request: web.Request) -> web.WebSocketResponse:
@@ -72,6 +73,10 @@ async def handler(self, request: web.Request) -> web.WebSocketResponse:
                 data = msg.json()
                 if self.next_on_subscribe and data.get("type") == "subscribe":
                     await ws.send_json({"id": data["id"], "type": "next"})
+                if self.error_on_subscribe and data.get("type") == "subscribe":
+                    await ws.send_json(
+                        {"id": data["id"], "type": "error", "payload": [{}]}
+                    )
         return ws
 
 
@@ -159,7 +164,7 @@ async def fake_sleep(delay: float) -> None:
 
 def add_subscription(monitor: WebSocketMonitor) -> None:
     """Give the monitor a subscription to resume, so it keeps reconnecting."""
-    monitor._subscriptions["sub"] = (lambda _: None, {"query": "subscription"})
+    monitor._subscriptions["sub"] = (lambda _: None, {"query": "subscription"}, None)
 
 
 async def wait_for(condition: Any, timeout: float = 5) -> None:
@@ -580,3 +585,18 @@ async def connection_init(ws: ClientWebSocketResponse[bool]) -> None:
             assert unacked.rate_limited
         finally:
             WebSocketMonitor._shared_backoffs.clear()
+
+
+async def test_subscription_errors_reach_on_error(
+    server: FakeServer, monitor: WebSocketMonitor
+) -> None:
+    """A subscription's error message goes to its on_error, if it has one."""
+    server.error_on_subscribe = True
+    await monitor.new_connection(start_monitor=True)
+    await wait_for(lambda: monitor.connection_ack.is_set())
+    errors: list[dict[str, Any]] = []
+    await monitor.start_subscription({"query": "a"}, lambda _: None, errors.append)
+    await monitor.start_subscription({"query": "b"}, lambda _: None)
+    await wait_for(lambda: errors)
+    assert errors[0]["type"] == "error"
+    assert monitor.connected