diff --git a/src/rivian/rivian.py b/src/rivian/rivian.py index 0d0ebc6..ea2878f 100644 --- a/src/rivian/rivian.py +++ b/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 diff --git a/src/rivian/ws_monitor.py b/src/rivian/ws_monitor.py index ccf0bce..c65dd1f 100644 --- a/src/rivian/ws_monitor.py +++ b/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: diff --git a/tests/test_ws_monitor.py b/tests/test_ws_monitor.py index 3e9909b..dd5da05 100644 --- a/tests/test_ws_monitor.py +++ b/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