chore: Apply co-pilot feedback

Author: allenporter

roborock/devices/rpc/v1_channel.py

@@ -308,14 +308,17 @@ async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callab
             loop = asyncio.get_running_loop()
             self._reconnect_task = loop.create_task(self._background_reconnect())
 
-        # Always subscribe to MQTT to receive protocol updates (data points)
+        # Always attempt to subscribe to MQTT to receive protocol updates (data points)
         # even if we have a local connection. Protocol updates only come via cloud/MQTT.
         # Local connection is used for RPC commands, but push notifications come via MQTT.
-        self._mqtt_unsub = await self._mqtt_channel.subscribe(self._on_mqtt_message)
-        if self.is_local_connected:
-            self._logger.debug("V1Channel connected via local and MQTT (for protocol updates)")
-        else:
-            self._logger.debug("V1Channel connected via MQTT only")
+        try:
+            self._mqtt_unsub = await self._mqtt_channel.subscribe(self._on_mqtt_message)
+        except RoborockException as err:
+            if not self.is_local_connected:
+                # Propagate error if both local and MQTT failed
+                self._logger.debug("MQTT connection also failed: %s", err)
+                raise
+            self._logger.debug("MQTT subscription failed, continuing with local-only connection: %s", err)
 
         def unsub() -> None:
             """Unsubscribe from all messages."""

roborock/devices/traits/v1/__init__.py

@@ -235,12 +235,15 @@ def _get_rpc_channel(self, trait: V1TraitMixin) -> V1RpcChannel:
 
     async def start(self) -> None:
         """Start the properties API and discover features."""
+        if self._unsub:
+            return
         await self.discover_features()
         self._unsub = self._add_dps_listener(self._on_dps_update)
 
     def close(self) -> None:
         if self._unsub:
             self._unsub()
+            self._unsub = None
 
     def _on_dps_update(self, dps: dict[RoborockDataProtocol, Any]) -> None:
         """Handle incoming messages from the device.

roborock/devices/traits/v1/consumeable.py

@@ -4,6 +4,7 @@
 periodically, such as filters, brushes, etc.
 """
 
+import logging
 from enum import StrEnum
 from typing import Any, Self
 
@@ -18,6 +19,8 @@
     "ConsumableTrait",
 ]
 
+_LOGGER = logging.getLogger(__name__)
+
 _DPS_CONVERTER = common.DpsDataConverter.from_dataclass(Consumable)
 
 
@@ -48,6 +51,11 @@ class ConsumableTrait(Consumable, common.V1TraitMixin, TraitUpdateListener):
     command = RoborockCommand.GET_CONSUMABLE
     converter = common.DefaultConverter(Consumable)
 
+    def __init__(self) -> None:
+        """Initialize the consumable trait."""
+        super().__init__()
+        TraitUpdateListener.__init__(self, logger=_LOGGER)
+
     async def reset_consumable(self, consumable: ConsumableAttribute) -> None:
         """Reset a specific consumable attribute on the device."""
         await self.rpc_channel.send_command(RoborockCommand.RESET_CONSUMABLE, params=[consumable.value])

roborock/protocols/v1_protocol.py

@@ -170,9 +170,14 @@ def decode_rpc_response(message: RoborockMessage) -> ResponseMessage:
     response, as long as we can extract the request ID. This is so we can
     associate an API response with a request even if there was an error.
     """
-    if not (datapoints := _decode_dps_message(message)):
+    if not message.payload:
         return ResponseMessage(request_id=message.seq, data={})
 
+    if (datapoints := _decode_dps_message(message)) is None:
+        raise RoborockException(
+            f"Invalid V1 message format: missing or invalid 'dps' in payload for {message.payload!r}"
+        )
+
     if not (data_point := datapoints.get(RoborockMessageProtocol.RPC_RESPONSE)):
         raise RoborockException(
             f"Invalid V1 message format: missing '{RoborockMessageProtocol.RPC_RESPONSE}' data point"

tests/devices/rpc/test_v1_channel.py

@@ -23,7 +23,7 @@
     create_mqtt_encoder,
 )
 from roborock.protocols.v1_protocol import MapResponse, SecurityData, V1RpcChannel
-from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
+from roborock.roborock_message import RoborockDataProtocol, RoborockMessage, RoborockMessageProtocol
 from roborock.roborock_typing import RoborockCommand
 from tests import mock_data
 from tests.fixtures.channel_fixtures import FakeChannel
@@ -580,3 +580,35 @@ async def test_v1_channel_send_map_command(
 
     # Verify the result is the data from our mocked decoder
     assert result == decompressed_map_data
+
+
+async def test_v1_channel_add_dps_listener(
+    v1_channel: V1Channel,
+    mock_mqtt_channel: FakeChannel,
+) -> None:
+    """Test that DPS listeners receive decoded protocol updates from MQTT."""
+    mock_mqtt_channel.response_queue.append(TEST_NETWORK_INFO_RESPONSE)
+    await v1_channel.subscribe(Mock())
+
+    # Create a mock listener for DPS updates
+    dps_listener = Mock()
+    unsub_dps = v1_channel.add_dps_listener(dps_listener)
+
+    # Simulate an incoming MQTT message with data protocol payload.
+    dps_payload = json.dumps({"dps": {"121": 5}}).encode()
+    push_message = RoborockMessage(
+        protocol=RoborockMessageProtocol.GENERAL_REQUEST,
+        payload=dps_payload,
+    )
+    mock_mqtt_channel.notify_subscribers(push_message)
+
+    dps_listener.assert_called_once()
+    called_args = dps_listener.call_args[0][0]
+    assert called_args[RoborockDataProtocol.STATE] == 5
+
+    unsub_dps()
+
+    # Verify unsubscribe works
+    dps_listener.reset_mock()
+    v1_channel._on_mqtt_message(push_message)
+    dps_listener.assert_not_called()

tests/fixtures/channel_fixtures.py

@@ -51,3 +51,11 @@ async def _subscribe(self, callback: Callable[[RoborockMessage], None]) -> Calla
         """Simulate subscribing to messages."""
         self.subscribers.append(callback)
         return lambda: self.subscribers.remove(callback)
+
+    def notify_subscribers(self, message: RoborockMessage) -> None:
+        """Notify subscribers of a message.
+
+        This can be used by tests to simulate the channel receiving a message.
+        """
+        for subscriber in list(self.subscribers):
+            subscriber(message)