chore: Move a01 encoding and decoding to a separate module

allenporter · allenporter · commit 054b3fb4b044 · 2025-07-16T14:25:07.000-07:00
diff --git a/roborock/protocols/a01_protocol.py b/roborock/protocols/a01_protocol.py
@@ -0,0 +1,54 @@
+"""Roborock A01 Protocol encoding and decoding."""
+
+import json
+import logging
+from typing import Any
+
+from Crypto.Cipher import AES
+from Crypto.Util.Padding import pad, unpad
+
+from roborock.exceptions import RoborockException
+from roborock.roborock_message import (
+    RoborockDyadDataProtocol,
+    RoborockMessage,
+    RoborockMessageProtocol,
+    RoborockZeoProtocol,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+A01_VERSION = b"A01"
+
+
+def encode_mqtt_payload(data: dict[RoborockDyadDataProtocol | RoborockZeoProtocol, Any]) -> RoborockMessage:
+    """Encode payload for A01 commands over MQTT."""
+    dps_data = {"dps": data}
+    payload = pad(json.dumps(dps_data).encode("utf-8"), AES.block_size)
+    return RoborockMessage(
+        protocol=RoborockMessageProtocol.RPC_REQUEST,
+        version=A01_VERSION,
+        payload=payload,
+    )
+
+
+def decode_rpc_response(message: RoborockMessage) -> dict[int, Any]:
+    """Decode a V1 RPC_RESPONSE message."""
+    if not message.payload:
+        raise RoborockException("Invalid A01 message format: missing payload")
+    try:
+        unpadded = unpad(message.payload, AES.block_size)
+    except ValueError as err:
+        raise RoborockException(f"Unable to unpad A01 payload: {err}")
+
+    try:
+        payload = json.loads(unpadded.decode())
+    except (json.JSONDecodeError, TypeError) as e:
+        raise RoborockException(f"Invalid A01 message payload: {e} for {message.payload!r}") from e
+
+    datapoints = payload.get("dps", {})
+    if not isinstance(datapoints, dict):
+        raise RoborockException(f"Invalid A01 message format: 'dps' should be a dictionary for {message.payload!r}")
+    try:
+        return {int(key): value for key, value in datapoints.items()}
+    except ValueError:
+        raise RoborockException(f"Invalid A01 message format: 'dps' key should be an integer for {message.payload!r}")
diff --git a/roborock/version_a01_apis/roborock_client_a01.py b/roborock/version_a01_apis/roborock_client_a01.py
@@ -1,13 +1,8 @@
-import dataclasses
-import json
 import logging
-import typing
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from datetime import time
-
-from Crypto.Cipher import AES
-from Crypto.Util.Padding import unpad
+from typing import Any
 
 from roborock import DeviceData
 from roborock.api import RoborockClient
@@ -33,6 +28,8 @@
     ZeoTemperature,
 )
 from roborock.containers import DyadProductInfo, DyadSndState, RoborockCategory
+from roborock.exceptions import RoborockException
+from roborock.protocols.a01_protocol import decode_rpc_response
 from roborock.roborock_message import (
     RoborockDyadDataProtocol,
     RoborockMessage,
@@ -43,111 +40,120 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-@dataclasses.dataclass
-class A01ProtocolCacheEntry:
-    post_process_fn: Callable
-    value: typing.Any | None = None
-
-
 # Right now this cache is not active, it was too much complexity for the initial addition of dyad.
-protocol_entries = {
-    RoborockDyadDataProtocol.STATUS: A01ProtocolCacheEntry(lambda val: RoborockDyadStateCode(val).name),
-    RoborockDyadDataProtocol.SELF_CLEAN_MODE: A01ProtocolCacheEntry(lambda val: DyadSelfCleanMode(val).name),
-    RoborockDyadDataProtocol.SELF_CLEAN_LEVEL: A01ProtocolCacheEntry(lambda val: DyadSelfCleanLevel(val).name),
-    RoborockDyadDataProtocol.WARM_LEVEL: A01ProtocolCacheEntry(lambda val: DyadWarmLevel(val).name),
-    RoborockDyadDataProtocol.CLEAN_MODE: A01ProtocolCacheEntry(lambda val: DyadCleanMode(val).name),
-    RoborockDyadDataProtocol.SUCTION: A01ProtocolCacheEntry(lambda val: DyadSuction(val).name),
-    RoborockDyadDataProtocol.WATER_LEVEL: A01ProtocolCacheEntry(lambda val: DyadWaterLevel(val).name),
-    RoborockDyadDataProtocol.BRUSH_SPEED: A01ProtocolCacheEntry(lambda val: DyadBrushSpeed(val).name),
-    RoborockDyadDataProtocol.POWER: A01ProtocolCacheEntry(lambda val: int(val)),
-    RoborockDyadDataProtocol.AUTO_DRY: A01ProtocolCacheEntry(lambda val: bool(val)),
-    RoborockDyadDataProtocol.MESH_LEFT: A01ProtocolCacheEntry(lambda val: int(360000 - val * 60)),
-    RoborockDyadDataProtocol.BRUSH_LEFT: A01ProtocolCacheEntry(lambda val: int(360000 - val * 60)),
-    RoborockDyadDataProtocol.ERROR: A01ProtocolCacheEntry(lambda val: DyadError(val).name),
-    RoborockDyadDataProtocol.VOLUME_SET: A01ProtocolCacheEntry(lambda val: int(val)),
-    RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: A01ProtocolCacheEntry(lambda val: bool(val)),
-    RoborockDyadDataProtocol.AUTO_DRY_MODE: A01ProtocolCacheEntry(lambda val: bool(val)),
-    RoborockDyadDataProtocol.SILENT_DRY_DURATION: A01ProtocolCacheEntry(lambda val: int(val)),  # in minutes
-    RoborockDyadDataProtocol.SILENT_MODE: A01ProtocolCacheEntry(lambda val: bool(val)),
-    RoborockDyadDataProtocol.SILENT_MODE_START_TIME: A01ProtocolCacheEntry(
-        lambda val: time(hour=int(val / 60), minute=val % 60)
+DYAD_PROTOCOL_ENTRIES: dict[RoborockDyadDataProtocol, Callable] = {
+    RoborockDyadDataProtocol.STATUS: lambda val: RoborockDyadStateCode(val).name,
+    RoborockDyadDataProtocol.SELF_CLEAN_MODE: lambda val: DyadSelfCleanMode(val).name,
+    RoborockDyadDataProtocol.SELF_CLEAN_LEVEL: lambda val: DyadSelfCleanLevel(val).name,
+    RoborockDyadDataProtocol.WARM_LEVEL: lambda val: DyadWarmLevel(val).name,
+    RoborockDyadDataProtocol.CLEAN_MODE: lambda val: DyadCleanMode(val).name,
+    RoborockDyadDataProtocol.SUCTION: lambda val: DyadSuction(val).name,
+    RoborockDyadDataProtocol.WATER_LEVEL: lambda val: DyadWaterLevel(val).name,
+    RoborockDyadDataProtocol.BRUSH_SPEED: lambda val: DyadBrushSpeed(val).name,
+    RoborockDyadDataProtocol.POWER: lambda val: int(val),
+    RoborockDyadDataProtocol.AUTO_DRY: lambda val: bool(val),
+    RoborockDyadDataProtocol.MESH_LEFT: lambda val: int(360000 - val * 60),
+    RoborockDyadDataProtocol.BRUSH_LEFT: lambda val: int(360000 - val * 60),
+    RoborockDyadDataProtocol.ERROR: lambda val: DyadError(val).name,
+    RoborockDyadDataProtocol.VOLUME_SET: lambda val: int(val),
+    RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: lambda val: bool(val),
+    RoborockDyadDataProtocol.AUTO_DRY_MODE: lambda val: bool(val),
+    RoborockDyadDataProtocol.SILENT_DRY_DURATION: lambda val: int(val),  # in minutes
+    RoborockDyadDataProtocol.SILENT_MODE: lambda val: bool(val),
+    RoborockDyadDataProtocol.SILENT_MODE_START_TIME: lambda val: time(
+        hour=int(val / 60), minute=val % 60
     ),  # in minutes since 00:00
-    RoborockDyadDataProtocol.SILENT_MODE_END_TIME: A01ProtocolCacheEntry(
-        lambda val: time(hour=int(val / 60), minute=val % 60)
+    RoborockDyadDataProtocol.SILENT_MODE_END_TIME: lambda val: time(
+        hour=int(val / 60), minute=val % 60
     ),  # in minutes since 00:00
-    RoborockDyadDataProtocol.RECENT_RUN_TIME: A01ProtocolCacheEntry(
-        lambda val: [int(v) for v in val.split(",")]
-    ),  # minutes of cleaning in past few days.
-    RoborockDyadDataProtocol.TOTAL_RUN_TIME: A01ProtocolCacheEntry(lambda val: int(val)),
-    RoborockDyadDataProtocol.SND_STATE: A01ProtocolCacheEntry(lambda val: DyadSndState.from_dict(val)),
-    RoborockDyadDataProtocol.PRODUCT_INFO: A01ProtocolCacheEntry(lambda val: DyadProductInfo.from_dict(val)),
+    RoborockDyadDataProtocol.RECENT_RUN_TIME: lambda val: [
+        int(v) for v in val.split(",")
+    ],  # minutes of cleaning in past few days.
+    RoborockDyadDataProtocol.TOTAL_RUN_TIME: lambda val: int(val),
+    RoborockDyadDataProtocol.SND_STATE: lambda val: DyadSndState.from_dict(val),
+    RoborockDyadDataProtocol.PRODUCT_INFO: lambda val: DyadProductInfo.from_dict(val),
 }
 
-zeo_data_protocol_entries = {
+ZEO_PROTOCOL_ENTRIES: dict[RoborockZeoProtocol, Callable] = {
     # ro
-    RoborockZeoProtocol.STATE: A01ProtocolCacheEntry(lambda val: ZeoState(val).name),
-    RoborockZeoProtocol.COUNTDOWN: A01ProtocolCacheEntry(lambda val: int(val)),
-    RoborockZeoProtocol.WASHING_LEFT: A01ProtocolCacheEntry(lambda val: int(val)),
-    RoborockZeoProtocol.ERROR: A01ProtocolCacheEntry(lambda val: ZeoError(val).name),
-    RoborockZeoProtocol.TIMES_AFTER_CLEAN: A01ProtocolCacheEntry(lambda val: int(val)),
-    RoborockZeoProtocol.DETERGENT_EMPTY: A01ProtocolCacheEntry(lambda val: bool(val)),
-    RoborockZeoProtocol.SOFTENER_EMPTY: A01ProtocolCacheEntry(lambda val: bool(val)),
+    RoborockZeoProtocol.STATE: lambda val: ZeoState(val).name,
+    RoborockZeoProtocol.COUNTDOWN: lambda val: int(val),
+    RoborockZeoProtocol.WASHING_LEFT: lambda val: int(val),
+    RoborockZeoProtocol.ERROR: lambda val: ZeoError(val).name,
+    RoborockZeoProtocol.TIMES_AFTER_CLEAN: lambda val: int(val),
+    RoborockZeoProtocol.DETERGENT_EMPTY: lambda val: bool(val),
+    RoborockZeoProtocol.SOFTENER_EMPTY: lambda val: bool(val),
     # rw
-    RoborockZeoProtocol.MODE: A01ProtocolCacheEntry(lambda val: ZeoMode(val).name),
-    RoborockZeoProtocol.PROGRAM: A01ProtocolCacheEntry(lambda val: ZeoProgram(val).name),
-    RoborockZeoProtocol.TEMP: A01ProtocolCacheEntry(lambda val: ZeoTemperature(val).name),
-    RoborockZeoProtocol.RINSE_TIMES: A01ProtocolCacheEntry(lambda val: ZeoRinse(val).name),
-    RoborockZeoProtocol.SPIN_LEVEL: A01ProtocolCacheEntry(lambda val: ZeoSpin(val).name),
-    RoborockZeoProtocol.DRYING_MODE: A01ProtocolCacheEntry(lambda val: ZeoDryingMode(val).name),
-    RoborockZeoProtocol.DETERGENT_TYPE: A01ProtocolCacheEntry(lambda val: ZeoDetergentType(val).name),
-    RoborockZeoProtocol.SOFTENER_TYPE: A01ProtocolCacheEntry(lambda val: ZeoSoftenerType(val).name),
-    RoborockZeoProtocol.SOUND_SET: A01ProtocolCacheEntry(lambda val: bool(val)),
+    RoborockZeoProtocol.MODE: lambda val: ZeoMode(val).name,
+    RoborockZeoProtocol.PROGRAM: lambda val: ZeoProgram(val).name,
+    RoborockZeoProtocol.TEMP: lambda val: ZeoTemperature(val).name,
+    RoborockZeoProtocol.RINSE_TIMES: lambda val: ZeoRinse(val).name,
+    RoborockZeoProtocol.SPIN_LEVEL: lambda val: ZeoSpin(val).name,
+    RoborockZeoProtocol.DRYING_MODE: lambda val: ZeoDryingMode(val).name,
+    RoborockZeoProtocol.DETERGENT_TYPE: lambda val: ZeoDetergentType(val).name,
+    RoborockZeoProtocol.SOFTENER_TYPE: lambda val: ZeoSoftenerType(val).name,
+    RoborockZeoProtocol.SOUND_SET: lambda val: bool(val),
 }
 
 
+def convert_dyad_value(protocol: int, value: Any) -> Any:
+    """Convert a dyad protocol value to its corresponding type."""
+    protocol_value = RoborockDyadDataProtocol(protocol)
+    if (converter := DYAD_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
+        return converter(value)
+    return None
+
+
+def convert_zeo_value(protocol: int, value: Any) -> Any:
+    """Convert a zeo protocol value to its corresponding type."""
+    protocol_value = RoborockZeoProtocol(protocol)
+    if (converter := ZEO_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
+        return converter(value)
+    return None
+
+
 class RoborockClientA01(RoborockClient, ABC):
     """Roborock client base class for A01 devices."""
 
+    value_converter: Callable[[int, Any], Any] | None = None
+
     def __init__(self, device_info: DeviceData, category: RoborockCategory):
         """Initialize the Roborock client."""
         super().__init__(device_info)
-        self.category = category
+        if category == RoborockCategory.WET_DRY_VAC:
+            self.value_converter = convert_dyad_value
+        elif category == RoborockCategory.WASHING_MACHINE:
+            self.value_converter = convert_zeo_value
+        else:
+            _LOGGER.debug("Device category %s is not (yet) supported", category)
+            self.value_converter = None
 
     def on_message_received(self, messages: list[RoborockMessage]) -> None:
+        if self.value_converter is None:
+            return
         for message in messages:
             protocol = message.protocol
             if message.payload and protocol in [
                 RoborockMessageProtocol.RPC_RESPONSE,
                 RoborockMessageProtocol.GENERAL_REQUEST,
             ]:
-                payload = message.payload
                 try:
-                    payload = unpad(payload, AES.block_size)
-                except Exception as err:
-                    self._logger.debug("Failed to unpad payload: %s", err)
+                    data_points = decode_rpc_response(message)
+                except RoborockException as err:
+                    self._logger.error("Failed to decode message %s: %s", message, err)
                     continue
-                payload_json = json.loads(payload.decode())
-                for data_point_number, data_point in payload_json.get("dps").items():
-                    data_point_protocol: RoborockDyadDataProtocol | RoborockZeoProtocol
-                    self._logger.debug("received msg with dps, protocol: %s, %s", data_point_number, protocol)
-                    entries: dict
-                    if self.category == RoborockCategory.WET_DRY_VAC:
-                        data_point_protocol = RoborockDyadDataProtocol(int(data_point_number))
-                        entries = protocol_entries
-                    elif self.category == RoborockCategory.WASHING_MACHINE:
-                        data_point_protocol = RoborockZeoProtocol(int(data_point_number))
-                        entries = zeo_data_protocol_entries
-                    else:
-                        continue
-                    if data_point_protocol in entries:
-                        # Auto convert into data struct we want.
-                        converted_response = entries[data_point_protocol].post_process_fn(data_point)
+                for data_point_number, data_point in data_points.items():
+                    if converted_response := self.value_converter(data_point_number, data_point):
                         queue = self._waiting_queue.get(int(data_point_number))
                         if queue and queue.protocol == protocol:
                             queue.set_result(converted_response)
+                    else:
+                        self._logger.warning(
+                            "Received unknown data point %s for protocol %s, ignoring", data_point_number, protocol
+                        )
 
     @abstractmethod
     async def update_values(
         self, dyad_data_protocols: list[RoborockDyadDataProtocol | RoborockZeoProtocol]
-    ) -> dict[RoborockDyadDataProtocol | RoborockZeoProtocol, typing.Any]:
+    ) -> dict[RoborockDyadDataProtocol | RoborockZeoProtocol, Any]:
         """This should handle updating for each given protocol."""
diff --git a/roborock/version_a01_apis/roborock_mqtt_client_a01.py b/roborock/version_a01_apis/roborock_mqtt_client_a01.py
@@ -4,11 +4,12 @@
 import typing
 
 from Crypto.Cipher import AES
-from Crypto.Util.Padding import pad, unpad
+from Crypto.Util.Padding import unpad
 
 from roborock.cloud_api import RoborockMqttClient
 from roborock.containers import DeviceData, RoborockCategory, UserData
 from roborock.exceptions import RoborockException
+from roborock.protocols.a01_protocol import encode_mqtt_payload
 from roborock.roborock_message import (
     RoborockDyadDataProtocol,
     RoborockMessage,
@@ -46,11 +47,13 @@ async def send_message(self, roborock_message: RoborockMessage):
         # self._logger.debug(f"id={request_id} Requesting method {method} with {params}")
         payload = json.loads(unpad(roborock_message.payload, AES.block_size))
         futures = []
+        self._logger.debug("Sending message: %s", payload)
         if "10000" in payload["dps"]:
             for dps in json.loads(payload["dps"]["10000"]):
                 futures.append(self._async_response(dps, response_protocol))
         self._send_msg_raw(m)
         responses = await asyncio.gather(*futures, return_exceptions=True)
+        self._logger.debug("Received responses: %s", responses)
         dps_responses: dict[int, typing.Any] = {}
         if "10000" in payload["dps"]:
             for i, dps in enumerate(json.loads(payload["dps"]["10000"])):
@@ -65,24 +68,14 @@ async def send_message(self, roborock_message: RoborockMessage):
     async def update_values(
         self, dyad_data_protocols: list[RoborockDyadDataProtocol | RoborockZeoProtocol]
     ) -> dict[RoborockDyadDataProtocol | RoborockZeoProtocol, typing.Any]:
-        payload = {"dps": {RoborockDyadDataProtocol.ID_QUERY: str([int(protocol) for protocol in dyad_data_protocols])}}
-        return await self.send_message(
-            RoborockMessage(
-                protocol=RoborockMessageProtocol.RPC_REQUEST,
-                version=b"A01",
-                payload=pad(json.dumps(payload).encode("utf-8"), AES.block_size),
-            )
+        message = encode_mqtt_payload(
+            {RoborockDyadDataProtocol.ID_QUERY: str([int(protocol) for protocol in dyad_data_protocols])}
         )
+        return await self.send_message(message)
 
     async def set_value(
         self, protocol: RoborockDyadDataProtocol | RoborockZeoProtocol, value: typing.Any
     ) -> dict[int, typing.Any]:
         """Set a value for a specific protocol on the A01 device."""
-        payload = {"dps": {int(protocol): value}}
-        return await self.send_message(
-            RoborockMessage(
-                protocol=RoborockMessageProtocol.RPC_REQUEST,
-                version=b"A01",
-                payload=pad(json.dumps(payload).encode("utf-8"), AES.block_size),
-            )
-        )
+        message = encode_mqtt_payload({protocol: value})
+        return await self.send_message(message)
diff --git a/tests/protocols/__init__.py b/tests/protocols/__init__.py
@@ -0,0 +1 @@
+"""Tests for the protocols package."""
diff --git a/tests/protocols/test_a01_protocol.py b/tests/protocols/test_a01_protocol.py
diff --git a/tests/test_a01_api.py b/tests/test_a01_api.py
