feat(q10): add decoded command helper and status trait

Author: lboue

roborock/devices/rpc/b01_q10_channel.py

@@ -2,17 +2,23 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
+from collections.abc import Iterable
+from typing import Any
 
 from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
 from roborock.devices.transport.mqtt_channel import MqttChannel
 from roborock.exceptions import RoborockException
 from roborock.protocols.b01_q10_protocol import (
     ParamsType,
+    decode_rpc_response,
     encode_mqtt_payload,
 )
+from roborock.roborock_message import RoborockMessage
 
 _LOGGER = logging.getLogger(__name__)
+_TIMEOUT = 10.0
 
 
 async def send_command(
@@ -34,3 +40,61 @@ async def send_command(
             ex,
         )
         raise
+
+
+async def send_decoded_command(
+    mqtt_channel: MqttChannel,
+    command: B01_Q10_DP,
+    params: ParamsType,
+    expected_dps: Iterable[B01_Q10_DP] | None = None,
+) -> dict[B01_Q10_DP, Any]:
+    """Send a command and await the first decoded response.
+
+    Q10 responses are not correlated with a message id, so we filter on
+    expected datapoints when provided.
+    """
+    roborock_message = encode_mqtt_payload(command, params)
+    future: asyncio.Future[dict[B01_Q10_DP, Any]] = asyncio.get_running_loop().create_future()
+
+    expected_set = set(expected_dps) if expected_dps is not None else None
+
+    def find_response(response_message: RoborockMessage) -> None:
+        try:
+            decoded_dps = decode_rpc_response(response_message)
+        except RoborockException as ex:
+            _LOGGER.debug(
+                "Failed to decode B01 Q10 RPC response (expecting %s): %s: %s",
+                command,
+                response_message,
+                ex,
+            )
+            return
+        if expected_set and not any(dps in decoded_dps for dps in expected_set):
+            return
+        if not future.done():
+            future.set_result(decoded_dps)
+
+    unsub = await mqtt_channel.subscribe(find_response)
+
+    _LOGGER.debug("Sending MQTT message: %s", roborock_message)
+    try:
+        await mqtt_channel.publish(roborock_message)
+        return await asyncio.wait_for(future, timeout=_TIMEOUT)
+    except TimeoutError as ex:
+        raise RoborockException(f"B01 Q10 command timed out after {_TIMEOUT}s ({command})") from ex
+    except RoborockException as ex:
+        _LOGGER.warning(
+            "Error sending B01 Q10 decoded command (%s): %s",
+            command,
+            ex,
+        )
+        raise
+    except Exception as ex:
+        _LOGGER.exception(
+            "Error sending B01 Q10 decoded command (%s): %s",
+            command,
+            ex,
+        )
+        raise
+    finally:
+        unsub()

roborock/devices/traits/b01/q10/status.py

@@ -0,0 +1,73 @@
+"""Status trait for Q10 B01 devices."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+from roborock.data.b01_q10.b01_q10_code_mappings import (
+    B01_Q10_DP,
+    YXDeviceCleanTask,
+    YXDeviceState,
+    YXDeviceWorkMode,
+    YXFanLevel,
+)
+from roborock.devices.rpc.b01_q10_channel import send_decoded_command
+from roborock.devices.transport.mqtt_channel import MqttChannel
+
+
+class StatusTrait:
+    """Trait for requesting and holding Q10 status values."""
+
+    def __init__(self, channel: MqttChannel) -> None:
+        self._channel = channel
+        self._data: dict[B01_Q10_DP, Any] = {}
+
+    @property
+    def data(self) -> dict[B01_Q10_DP, Any]:
+        """Return the latest raw status data."""
+        return self._data
+
+    async def refresh(self) -> dict[B01_Q10_DP, Any]:
+        """Refresh status values from the device."""
+        decoded = await send_decoded_command(
+            self._channel,
+            command=B01_Q10_DP.REQUEST_DPS,
+            params={},
+            expected_dps={B01_Q10_DP.STATUS, B01_Q10_DP.BATTERY},
+        )
+        self._data = decoded
+        return decoded
+
+    @property
+    def state_code(self) -> int | None:
+        return self._data.get(B01_Q10_DP.STATUS)
+
+    @property
+    def state(self) -> YXDeviceState | None:
+        code = self.state_code
+        return cast(YXDeviceState | None, YXDeviceState.from_code_optional(code)) if code is not None else None
+
+    @property
+    def battery(self) -> int | None:
+        return self._data.get(B01_Q10_DP.BATTERY)
+
+    @property
+    def fan_level(self) -> YXFanLevel | None:
+        value = self._data.get(B01_Q10_DP.FAN_LEVEL)
+        return cast(YXFanLevel | None, YXFanLevel.from_code_optional(value)) if value is not None else None
+
+    @property
+    def clean_mode(self) -> YXDeviceWorkMode | None:
+        value = self._data.get(B01_Q10_DP.CLEAN_MODE)
+        return cast(YXDeviceWorkMode | None, YXDeviceWorkMode.from_code_optional(value)) if value is not None else None
+
+    @property
+    def clean_task(self) -> YXDeviceCleanTask | None:
+        value = self._data.get(B01_Q10_DP.CLEAN_TASK_TYPE)
+        return (
+            cast(YXDeviceCleanTask | None, YXDeviceCleanTask.from_code_optional(value)) if value is not None else None
+        )
+
+    @property
+    def cleaning_progress(self) -> int | None:
+        return self._data.get(B01_Q10_DP.CLEANING_PROGRESS)