fix: convert a01 values

roborock/devices/traits/a01/__init__.py

@@ -1,6 +1,31 @@
+from collections.abc import Callable
+from datetime import time
 from typing import Any
 
-from roborock.data import HomeDataProduct, RoborockCategory
+from roborock.data import DyadProductInfo, DyadSndState, HomeDataProduct, RoborockCategory
+from roborock.data.dyad.dyad_code_mappings import (
+    DyadBrushSpeed,
+    DyadCleanMode,
+    DyadError,
+    DyadSelfCleanLevel,
+    DyadSelfCleanMode,
+    DyadSuction,
+    DyadWarmLevel,
+    DyadWaterLevel,
+    RoborockDyadStateCode,
+)
+from roborock.data.zeo.zeo_code_mappings import (
+    ZeoDetergentType,
+    ZeoDryingMode,
+    ZeoError,
+    ZeoMode,
+    ZeoProgram,
+    ZeoRinse,
+    ZeoSoftenerType,
+    ZeoSpin,
+    ZeoState,
+    ZeoTemperature,
+)
 from roborock.devices.a01_channel import send_decoded_command
 from roborock.devices.mqtt_channel import MqttChannel
 from roborock.devices.traits import Trait
@@ -12,6 +37,81 @@
 ]
 
 
+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: lambda val: time(
+        hour=int(val / 60), minute=val % 60
+    ),  # in minutes since 00:00
+    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_PROTOCOL_ENTRIES: dict[RoborockZeoProtocol, Callable] = {
+    # read-only
+    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),
+    # read-write
+    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_value: RoborockDyadDataProtocol, value: Any) -> Any:
+    """Convert a dyad protocol value to its corresponding type."""
+    if (converter := DYAD_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
+        try:
+            return converter(value)
+        except (ValueError, TypeError):
+            return None
+    return None
+
+
+def convert_zeo_value(protocol_value: RoborockZeoProtocol, value: Any) -> Any:
+    """Convert a zeo protocol value to its corresponding type."""
+    if (converter := ZEO_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
+        try:
+            return converter(value)
+        except (ValueError, TypeError):
+            return None
+    return None
+
+
 class DyadApi(Trait):
     """API for interacting with Dyad devices."""
 
@@ -22,7 +122,8 @@ def __init__(self, channel: MqttChannel) -> None:
     async def query_values(self, protocols: list[RoborockDyadDataProtocol]) -> dict[RoborockDyadDataProtocol, Any]:
         """Query the device for the values of the given Dyad protocols."""
         params = {RoborockDyadDataProtocol.ID_QUERY: str([int(p) for p in protocols])}
-        return await send_decoded_command(self._channel, params)
+        response = await send_decoded_command(self._channel, params)
+        return {protocol: convert_dyad_value(protocol, response.get(protocol)) for protocol in protocols}
 
     async def set_value(self, protocol: RoborockDyadDataProtocol, value: Any) -> dict[RoborockDyadDataProtocol, Any]:
         """Set a value for a specific protocol on the device."""
@@ -42,7 +143,8 @@ def __init__(self, channel: MqttChannel) -> None:
     async def query_values(self, protocols: list[RoborockZeoProtocol]) -> dict[RoborockZeoProtocol, Any]:
         """Query the device for the values of the given protocols."""
         params = {RoborockZeoProtocol.ID_QUERY: str([int(p) for p in protocols])}
-        return await send_decoded_command(self._channel, params)
+        response = await send_decoded_command(self._channel, params)
+        return {protocol: convert_zeo_value(protocol, response.get(protocol)) for protocol in protocols}
 
     async def set_value(self, protocol: RoborockZeoProtocol, value: Any) -> dict[RoborockZeoProtocol, Any]:
         """Set a value for a specific protocol on the device."""

tests/devices/test_a01_traits.py

@@ -0,0 +1,55 @@
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from roborock.devices.traits.a01 import DyadApi, ZeoApi
+from roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol
+
+
+@pytest.fixture
+def mock_channel():
+    channel = Mock()
+    channel.send_command = AsyncMock()
+    # Mocking send_decoded_command if it was a method on channel, but it's a standalone function imported in traits.
+    # However, in traits/__init__.py it is imported as: from roborock.devices.a01_channel import send_decoded_command
+    return channel
+
+
+@pytest.mark.asyncio
+async def test_dyad_query_values(mock_channel):
+    with patch("roborock.devices.traits.a01.send_decoded_command", new_callable=AsyncMock) as mock_send:
+        api = DyadApi(mock_channel)
+
+        # Setup mock return value (raw values)
+        mock_send.return_value = {
+            int(RoborockDyadDataProtocol.CLEAN_MODE): 1,
+            int(RoborockDyadDataProtocol.POWER): 100,
+        }
+
+        protocols = [RoborockDyadDataProtocol.CLEAN_MODE, RoborockDyadDataProtocol.POWER]
+        result = await api.query_values(protocols)
+
+        # Verify conversion
+        assert RoborockDyadDataProtocol.CLEAN_MODE in result
+        assert RoborockDyadDataProtocol.POWER in result
+
+        assert isinstance(result[RoborockDyadDataProtocol.CLEAN_MODE], str)
+        assert result[RoborockDyadDataProtocol.POWER] == 100
+
+
+@pytest.mark.asyncio
+async def test_zeo_query_values(mock_channel):
+    with patch("roborock.devices.traits.a01.send_decoded_command", new_callable=AsyncMock) as mock_send:
+        api = ZeoApi(mock_channel)
+
+        mock_send.return_value = {
+            int(RoborockZeoProtocol.STATE): 6,  # spinning
+            int(RoborockZeoProtocol.COUNTDOWN): 120,
+        }
+
+        protocols = [RoborockZeoProtocol.STATE, RoborockZeoProtocol.COUNTDOWN]
+        result = await api.query_values(protocols)
+
+        assert RoborockZeoProtocol.STATE in result
+        assert result[RoborockZeoProtocol.STATE] == "spinning"
+        assert result[RoborockZeoProtocol.COUNTDOWN] == 120

tests/devices/traits/a01/test_init.py

@@ -1,3 +1,4 @@
+import datetime
 from collections.abc import Generator
 from typing import Any
 from unittest.mock import AsyncMock, call, patch
@@ -51,17 +52,16 @@ async def test_dyad_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMo
         ]
     )
     assert result == {
-        # Note: Bugs here, returning raw values
         RoborockDyadDataProtocol.POWER: 1,
-        RoborockDyadDataProtocol.STATUS: 6,
-        RoborockDyadDataProtocol.WATER_LEVEL: 3,
-        RoborockDyadDataProtocol.MESH_LEFT: 120,
-        RoborockDyadDataProtocol.BRUSH_LEFT: 90,
-        RoborockDyadDataProtocol.SILENT_MODE_START_TIME: 85,
-        RoborockDyadDataProtocol.RECENT_RUN_TIME: "3,4,5",
+        RoborockDyadDataProtocol.STATUS: "self_clean_deep_cleaning",
+        RoborockDyadDataProtocol.WATER_LEVEL: "l3",
+        RoborockDyadDataProtocol.MESH_LEFT: 352800,
+        RoborockDyadDataProtocol.BRUSH_LEFT: 354600,
+        RoborockDyadDataProtocol.SILENT_MODE_START_TIME: datetime.time(1, 25),
+        RoborockDyadDataProtocol.RECENT_RUN_TIME: [3, 4, 5],
         RoborockDyadDataProtocol.TOTAL_RUN_TIME: 123456,
-        RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: 1,
-        RoborockDyadDataProtocol.AUTO_DRY_MODE: 0,
+        RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: True,
+        RoborockDyadDataProtocol.AUTO_DRY_MODE: False,
     }
 
     # Note: Bug here, this is the wrong encoding for the query
@@ -86,11 +86,7 @@ async def test_dyad_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMo
                 9999: -3,
             },
             {
-                # Note: Bug here, should return enum value
-                RoborockDyadDataProtocol.STATUS: 3,
-                # Note: Bug here, unknown value should not be returned
-                7: 1,
-                9999: -3,
+                RoborockDyadDataProtocol.STATUS: "charging",
             },
         ),
         (
@@ -99,8 +95,7 @@ async def test_dyad_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMo
                 RoborockDyadDataProtocol.SILENT_MODE_START_TIME: "invalid",
             },
             {
-                # Note: Bug here, invalid value should not be returned
-                RoborockDyadDataProtocol.SILENT_MODE_START_TIME: "invalid",
+                RoborockDyadDataProtocol.SILENT_MODE_START_TIME: None,
             },
         ),
         (
@@ -111,11 +106,7 @@ async def test_dyad_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMo
                 9999: -3,
             },
             {
-                # Note: Bug here, should return time value
-                RoborockDyadDataProtocol.SILENT_MODE_START_TIME: 85,
-                # Note: Bug here, additional values should not be returned
-                RoborockDyadDataProtocol.POWER: 2,
-                9999: -3,
+                RoborockDyadDataProtocol.SILENT_MODE_START_TIME: datetime.time(1, 25),
             },
         ),
     ],
@@ -164,10 +155,10 @@ async def test_zeo_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMoc
     )
     assert result == {
         # Note: Bug here, should return enum/bool values
-        RoborockZeoProtocol.STATE: 6,
-        RoborockZeoProtocol.TEMP: 3,
-        RoborockZeoProtocol.DETERGENT_EMPTY: 1,
-        RoborockZeoProtocol.SOFTENER_EMPTY: 0,
+        RoborockZeoProtocol.STATE: "spinning",
+        RoborockZeoProtocol.TEMP: "medium",
+        RoborockZeoProtocol.DETERGENT_EMPTY: True,
+        RoborockZeoProtocol.SOFTENER_EMPTY: False,
         RoborockZeoProtocol.TIMES_AFTER_CLEAN: 1,
         RoborockZeoProtocol.WASHING_LEFT: 0,
     }
@@ -193,11 +184,7 @@ async def test_zeo_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMoc
                 9999: -3,
             },
             {
-                # Note: Bug here, should return enum/bool value
-                RoborockZeoProtocol.STATE: 1,
-                # Note: Bug here, unknown value should not be returned
-                7: 1,
-                9999: -3,
+                RoborockZeoProtocol.STATE: "standby",
             },
         ),
         (
@@ -206,8 +193,7 @@ async def test_zeo_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMoc
                 RoborockZeoProtocol.WASHING_LEFT: "invalid",
             },
             {
-                # Note: Bug here, invalid value should not be returned
-                RoborockZeoProtocol.WASHING_LEFT: "invalid",
+                RoborockZeoProtocol.WASHING_LEFT: None,
             },
         ),
         (
@@ -218,10 +204,7 @@ async def test_zeo_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMoc
                 9999: -3,
             },
             {
-                RoborockZeoProtocol.STATE: 1,
-                # Note: Bug here, these values were not requested and should not be returned
-                RoborockZeoProtocol.WASHING_LEFT: 2,
-                9999: -3,
+                RoborockZeoProtocol.STATE: "standby",
             },
         ),
     ],