feat: Add a DnD trait and fix bugs in the rpc channels

Author: allenporter

roborock/containers.py

@@ -134,8 +134,8 @@ def from_dict(cls, data: dict[str, Any]):
             return None
         field_types = {field.name: field.type for field in dataclasses.fields(cls)}
         result: dict[str, Any] = {}
-        for key, value in data.items():
-            key = _decamelize(key)
+        for orig_key, value in data.items():
+            key = _decamelize(orig_key)
             if (field_type := field_types.get(key)) is None:
                 continue
             if value == "None" or value is None:
@@ -178,16 +178,18 @@ class RoborockBaseTimer(RoborockBase):
     end_hour: int | None = None
     end_minute: int | None = None
     enabled: int | None = None
-    start_time: datetime.time | None = None
-    end_time: datetime.time | None = None
 
-    def __post_init__(self) -> None:
-        self.start_time = (
+    @property
+    def start_time(self) -> datetime.time | None:
+        return (
             datetime.time(hour=self.start_hour, minute=self.start_minute)
             if self.start_hour is not None and self.start_minute is not None
             else None
         )
-        self.end_time = (
+
+    @property
+    def end_time(self) -> datetime.time | None:
+        return (
             datetime.time(hour=self.end_hour, minute=self.end_minute)
             if self.end_hour is not None and self.end_minute is not None
             else None

roborock/devices/device_manager.py

@@ -22,6 +22,7 @@
 from .channel import Channel
 from .mqtt_channel import create_mqtt_channel
 from .traits.b01.props import B01PropsApi
+from .traits.dnd import DoNotDisturbTrait
 from .traits.dyad import DyadApi
 from .traits.status import StatusTrait
 from .traits.trait import Trait
@@ -152,6 +153,7 @@ def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> Roborock
             case DeviceVersion.V1:
                 channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
                 traits.append(StatusTrait(product, channel.rpc_channel))
+                traits.append(DoNotDisturbTrait(channel.rpc_channel))
             case DeviceVersion.A01:
                 mqtt_channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
                 match product.category:

roborock/devices/traits/dnd.py

@@ -0,0 +1,42 @@
+"""Module for Roborock V1 devices.
+
+This interface is experimental and subject to breaking changes without notice
+until the API is stable.
+"""
+
+import logging
+from collections.abc import Callable
+
+from roborock.containers import DnDTimer
+from roborock.devices.v1_rpc_channel import V1RpcChannel
+from roborock.roborock_typing import RoborockCommand
+
+from .trait import Trait
+
+_LOGGER = logging.getLogger(__name__)
+
+__all__ = [
+    "DoNotDisturbTrait",
+]
+
+
+class DoNotDisturbTrait(Trait):
+    """Trait for managing Do Not Disturb (DND) settings on Roborock devices."""
+
+    name = "do_not_disturb"
+
+    def __init__(self, rpc_channel: Callable[[], V1RpcChannel]) -> None:
+        """Initialize the DoNotDisturbTrait."""
+        self._rpc_channel = rpc_channel
+
+    async def get_dnd_timer(self) -> DnDTimer:
+        """Get the current Do Not Disturb (DND) timer settings of the device."""
+        return await self._rpc_channel().send_command(RoborockCommand.GET_DND_TIMER, response_type=DnDTimer)
+
+    async def set_dnd_timer(self, dnd_timer: DnDTimer) -> None:
+        """Set the Do Not Disturb (DND) timer settings of the device."""
+        await self._rpc_channel().send_command(RoborockCommand.SET_DND_TIMER, params=dnd_timer.as_dict())
+
+    async def clear_dnd_timer(self) -> None:
+        """Clear the Do Not Disturb (DND) timer settings of the device."""
+        await self._rpc_channel().send_command(RoborockCommand.CLOSE_DND_TIMER)

roborock/devices/traits/status.py

@@ -5,31 +5,32 @@
 """
 
 import logging
+from collections.abc import Callable
 
 from roborock.containers import (
     HomeDataProduct,
     ModelStatus,
     S7MaxVStatus,
     Status,
 )
+from roborock.devices.v1_rpc_channel import V1RpcChannel
 from roborock.roborock_typing import RoborockCommand
 
-from ..v1_rpc_channel import V1RpcChannel
 from .trait import Trait
 
 _LOGGER = logging.getLogger(__name__)
 
 __all__ = [
-    "Status",
+    "StatusTrait",
 ]
 
 
 class StatusTrait(Trait):
-    """Unified Roborock device class with automatic connection setup."""
+    """Trait for managing the status of Roborock devices."""
 
     name = "status"
 
-    def __init__(self, product_info: HomeDataProduct, rpc_channel: V1RpcChannel) -> None:
+    def __init__(self, product_info: HomeDataProduct, rpc_channel: Callable[[], V1RpcChannel]) -> None:
         """Initialize the StatusTrait."""
         self._product_info = product_info
         self._rpc_channel = rpc_channel
@@ -40,4 +41,4 @@ async def get_status(self) -> Status:
         This is a placeholder command and will likely be changed/moved in the future.
         """
         status_type: type[Status] = ModelStatus.get(self._product_info.model, S7MaxVStatus)
-        return await self._rpc_channel.send_command(RoborockCommand.GET_STATUS, response_type=status_type)
+        return await self._rpc_channel().send_command(RoborockCommand.GET_STATUS, response_type=status_type)

roborock/devices/v1_channel.py

@@ -81,9 +81,12 @@ def is_mqtt_connected(self) -> bool:
         """Return whether MQTT connection is available."""
         return self._mqtt_unsub is not None and self._mqtt_channel.is_connected
 
-    @property
     def rpc_channel(self) -> V1RpcChannel:
-        """Return the combined RPC channel prefers local with a fallback to MQTT."""
+        """Return the combined RPC channel prefers local with a fallback to MQTT.
+
+        This is dynamic based on the current connection status. That is, it may return
+        a different channel depending on whether local or MQTT is available.
+        """
         return self._combined_rpc_channel or self._mqtt_rpc_channel
 
     @property

roborock/devices/v1_rpc_channel.py

@@ -132,17 +132,21 @@ async def _send_raw_command(
         params: ParamsType = None,
     ) -> Any:
         """Send a command and return a parsed response RoborockBase type."""
-        _LOGGER.debug("Sending command (%s): %s, params=%s", self._name, method, params)
         request_message = RequestMessage(method, params=params)
+        _LOGGER.debug(
+            "Sending command (%s, request_id=%s): %s, params=%s", self._name, request_message.request_id, method, params
+        )
         message = self._payload_encoder(request_message)
 
         future: asyncio.Future[dict[str, Any]] = asyncio.Future()
 
         def find_response(response_message: RoborockMessage) -> None:
             try:
                 decoded = decode_rpc_response(response_message)
-            except RoborockException:
+            except RoborockException as ex:
+                _LOGGER.debug("Exception while decoding message (%s): %s", response_message, ex)
                 return
+            _LOGGER.debug("Received response (request_id=%s): %s", self._name, decoded.request_id)
             if decoded.request_id == request_message.request_id:
                 future.set_result(decoded.data)
 

roborock/protocols/v1_protocol.py

@@ -109,7 +109,7 @@ class ResponseMessage:
 def decode_rpc_response(message: RoborockMessage) -> ResponseMessage:
     """Decode a V1 RPC_RESPONSE message."""
     if not message.payload:
-        raise RoborockException("Invalid V1 message format: missing payload")
+        return ResponseMessage(request_id=message.seq, data={})
     try:
         payload = json.loads(message.payload.decode())
     except (json.JSONDecodeError, TypeError) as e:
@@ -141,6 +141,8 @@ def decode_rpc_response(message: RoborockMessage) -> ResponseMessage:
     _LOGGER.debug("Decoded V1 message result: %s", result)
     if isinstance(result, list) and result:
         result = result[0]
+    if isinstance(result, str) and result == "ok":
+        result = {}
     if not isinstance(result, dict):
         raise RoborockException(f"Invalid V1 message format: 'result' should be a dictionary for {message.payload!r}")
     return ResponseMessage(request_id=request_id, data=result)

tests/devices/test_v1_channel.py

@@ -254,7 +254,7 @@ async def test_v1_channel_send_command_local_preferred(
 
     # Send command
     mock_local_channel.response_queue.append(TEST_RESPONSE)
-    result = await v1_channel.rpc_channel.send_command(
+    result = await v1_channel.rpc_channel().send_command(
         RoborockCommand.CHANGE_SOUND_VOLUME,
         response_type=S5MaxStatus,
     )
@@ -280,7 +280,7 @@ async def test_v1_channel_send_command_local_fails(
 
     # Send command
     with pytest.raises(RoborockException, match="Local failed"):
-        await v1_channel.rpc_channel.send_command(
+        await v1_channel.rpc_channel().send_command(
             RoborockCommand.CHANGE_SOUND_VOLUME,
             response_type=S5MaxStatus,
         )
@@ -300,7 +300,7 @@ async def test_v1_channel_send_decoded_command_mqtt_only(
 
     # Send command
     mock_mqtt_channel.response_queue.append(TEST_RESPONSE)
-    result = await v1_channel.rpc_channel.send_command(
+    result = await v1_channel.rpc_channel().send_command(
         RoborockCommand.CHANGE_SOUND_VOLUME,
         response_type=S5MaxStatus,
     )
@@ -322,7 +322,7 @@ async def test_v1_channel_send_decoded_command_with_params(
     # Send command with params
     mock_local_channel.response_queue.append(TEST_RESPONSE)
     test_params = {"volume": 80}
-    await v1_channel.rpc_channel.send_command(
+    await v1_channel.rpc_channel().send_command(
         RoborockCommand.CHANGE_SOUND_VOLUME,
         response_type=S5MaxStatus,
         params=test_params,
@@ -444,7 +444,7 @@ async def test_v1_channel_command_encoding_validation(
 
     # Send local command and capture the request
     mock_local_channel.response_queue.append(TEST_RESPONSE_2)
-    await v1_channel.rpc_channel.send_command(RoborockCommand.CHANGE_SOUND_VOLUME, params={"volume": 50})
+    await v1_channel.rpc_channel().send_command(RoborockCommand.CHANGE_SOUND_VOLUME, params={"volume": 50})
     assert mock_local_channel.published_messages
     local_message = mock_local_channel.published_messages[0]
 
@@ -512,7 +512,7 @@ async def test_v1_channel_full_subscribe_and_command_flow(
 
     # Send a command (should use local)
     mock_local_channel.response_queue.append(TEST_RESPONSE)
-    result = await v1_channel.rpc_channel.send_command(
+    result = await v1_channel.rpc_channel().send_command(
         RoborockCommand.GET_STATUS,
         response_type=S5MaxStatus,
     )

tests/devices/test_v1_device.py

@@ -44,7 +44,7 @@ def traits_fixture(rpc_channel: AsyncMock) -> list[Trait]:
     return [
         StatusTrait(
             product_info=HOME_DATA.products[0],
-            rpc_channel=rpc_channel,
+            rpc_channel=lambda: rpc_channel,
         )
     ]
 

tests/devices/traits/__init__.py

@@ -0,0 +1 @@
+"""Tests for device traits."""