From 63d10d1f8b6630245251f62e822e5b32122f55cd Mon Sep 17 00:00:00 2001 From: Luke Date: Sun, 7 Dec 2025 12:10:50 -0500 Subject: [PATCH 1/4] fix: convert a01 values --- roborock/devices/traits/a01/__init__.py | 112 +++++++++++++++++++++++- tests/devices/test_a01_traits.py | 73 +++++++++++++++ 2 files changed, 182 insertions(+), 3 deletions(-) create mode 100644 tests/devices/test_a01_traits.py diff --git a/roborock/devices/traits/a01/__init__.py b/roborock/devices/traits/a01/__init__.py index e804f23c..e6ea3a92 100644 --- a/roborock/devices/traits/a01/__init__.py +++ b/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,77 @@ ] +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] = { + # ro + 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: 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 DyadApi(Trait): """API for interacting with Dyad devices.""" @@ -22,7 +118,12 @@ 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 { + RoborockDyadDataProtocol(int(k)): v + for k, val in response.items() + if (v := convert_dyad_value(int(k), val)) is not None + } 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,12 @@ 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 { + RoborockZeoProtocol(int(k)): v + for k, val in response.items() + if (v := convert_zeo_value(int(k), val)) is not None + } async def set_value(self, protocol: RoborockZeoProtocol, value: Any) -> dict[RoborockZeoProtocol, Any]: """Set a value for a specific protocol on the device.""" diff --git a/tests/devices/test_a01_traits.py b/tests/devices/test_a01_traits.py new file mode 100644 index 00000000..29155b04 --- /dev/null +++ b/tests/devices/test_a01_traits.py @@ -0,0 +1,73 @@ +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 + # Implementation detail: we need to mock send_decoded_command where it is used. + return channel + + +@pytest.fixture +def mock_send_decoded_command(): + with patch("roborock.devices.traits.a01.send_decoded_command", new_callable=AsyncMock) as mock: + yield mock + + +@pytest.mark.asyncio +async def test_dyad_query_values(mock_channel): + # We need to patch send_decoded_command in the module under test + 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, # Should convert to DyadCleanMode(1).name -> AUTO? Check mapping or enum + int(RoborockDyadDataProtocol.POWER): 100, + } + + protocols = [RoborockDyadDataProtocol.CLEAN_MODE, RoborockDyadDataProtocol.POWER] + result = await api.query_values(protocols) + + # Verify conversion + # CLEAN_MODE 1 -> str + # POWER 100 -> 100 + + assert RoborockDyadDataProtocol.CLEAN_MODE in result + assert RoborockDyadDataProtocol.POWER in result + + # Check actual values if we know the mapping. + # From roborock_client_a01.py (now a01_conversions.py): + # RoborockDyadDataProtocol.CLEAN_MODE: lambda val: DyadCleanMode(val).name + # DyadCleanMode(1) would need to be checked. Let's just assert it is a string. + 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 + # From a01_conversions.py: RoborockZeoProtocol.STATE: lambda val: ZeoState(val).name + assert result[RoborockZeoProtocol.STATE] == "spinning" # Assuming ZeoState(6).name is spinning + assert result[RoborockZeoProtocol.COUNTDOWN] == 120 From d854ebad1a91569135925ed092ae830854388543 Mon Sep 17 00:00:00 2001 From: Luke Date: Sun, 7 Dec 2025 12:27:05 -0500 Subject: [PATCH 2/4] chore: small changes to comments --- roborock/devices/traits/a01/__init__.py | 4 ++-- tests/devices/test_a01_traits.py | 8 -------- 2 files changed, 2 insertions(+), 10 deletions(-) diff --git a/roborock/devices/traits/a01/__init__.py b/roborock/devices/traits/a01/__init__.py index e6ea3a92..53103f54 100644 --- a/roborock/devices/traits/a01/__init__.py +++ b/roborock/devices/traits/a01/__init__.py @@ -71,7 +71,7 @@ } ZEO_PROTOCOL_ENTRIES: dict[RoborockZeoProtocol, Callable] = { - # ro + # read-only RoborockZeoProtocol.STATE: lambda val: ZeoState(val).name, RoborockZeoProtocol.COUNTDOWN: lambda val: int(val), RoborockZeoProtocol.WASHING_LEFT: lambda val: int(val), @@ -79,7 +79,7 @@ RoborockZeoProtocol.TIMES_AFTER_CLEAN: lambda val: int(val), RoborockZeoProtocol.DETERGENT_EMPTY: lambda val: bool(val), RoborockZeoProtocol.SOFTENER_EMPTY: lambda val: bool(val), - # rw + # read-write RoborockZeoProtocol.MODE: lambda val: ZeoMode(val).name, RoborockZeoProtocol.PROGRAM: lambda val: ZeoProgram(val).name, RoborockZeoProtocol.TEMP: lambda val: ZeoTemperature(val).name, diff --git a/tests/devices/test_a01_traits.py b/tests/devices/test_a01_traits.py index 29155b04..7f35c7ea 100644 --- a/tests/devices/test_a01_traits.py +++ b/tests/devices/test_a01_traits.py @@ -12,19 +12,11 @@ def mock_channel(): 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 - # Implementation detail: we need to mock send_decoded_command where it is used. return channel -@pytest.fixture -def mock_send_decoded_command(): - with patch("roborock.devices.traits.a01.send_decoded_command", new_callable=AsyncMock) as mock: - yield mock - - @pytest.mark.asyncio async def test_dyad_query_values(mock_channel): - # We need to patch send_decoded_command in the module under test with patch("roborock.devices.traits.a01.send_decoded_command", new_callable=AsyncMock) as mock_send: api = DyadApi(mock_channel) From 3e85802dc9b39bf7bb3e77bf05b520b37614e651 Mon Sep 17 00:00:00 2001 From: Luke Date: Sun, 7 Dec 2025 13:23:27 -0500 Subject: [PATCH 3/4] chore: small changes to comments --- tests/devices/test_a01_traits.py | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/tests/devices/test_a01_traits.py b/tests/devices/test_a01_traits.py index 7f35c7ea..be400f0a 100644 --- a/tests/devices/test_a01_traits.py +++ b/tests/devices/test_a01_traits.py @@ -22,9 +22,7 @@ async def test_dyad_query_values(mock_channel): # Setup mock return value (raw values) mock_send.return_value = { - int( - RoborockDyadDataProtocol.CLEAN_MODE - ): 1, # Should convert to DyadCleanMode(1).name -> AUTO? Check mapping or enum + int(RoborockDyadDataProtocol.CLEAN_MODE): 1, int(RoborockDyadDataProtocol.POWER): 100, } @@ -32,16 +30,9 @@ async def test_dyad_query_values(mock_channel): result = await api.query_values(protocols) # Verify conversion - # CLEAN_MODE 1 -> str - # POWER 100 -> 100 - assert RoborockDyadDataProtocol.CLEAN_MODE in result assert RoborockDyadDataProtocol.POWER in result - # Check actual values if we know the mapping. - # From roborock_client_a01.py (now a01_conversions.py): - # RoborockDyadDataProtocol.CLEAN_MODE: lambda val: DyadCleanMode(val).name - # DyadCleanMode(1) would need to be checked. Let's just assert it is a string. assert isinstance(result[RoborockDyadDataProtocol.CLEAN_MODE], str) assert result[RoborockDyadDataProtocol.POWER] == 100 @@ -60,6 +51,5 @@ async def test_zeo_query_values(mock_channel): result = await api.query_values(protocols) assert RoborockZeoProtocol.STATE in result - # From a01_conversions.py: RoborockZeoProtocol.STATE: lambda val: ZeoState(val).name - assert result[RoborockZeoProtocol.STATE] == "spinning" # Assuming ZeoState(6).name is spinning + assert result[RoborockZeoProtocol.STATE] == "spinning" assert result[RoborockZeoProtocol.COUNTDOWN] == 120 From 97d8b7c43ba8291275bfcf5a30aa42365f68730a Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sun, 7 Dec 2025 10:53:27 -0800 Subject: [PATCH 4/4] fix: update tests and conversion logic --- roborock/devices/traits/a01/__init__.py | 28 ++++++------- tests/devices/traits/a01/test_init.py | 55 +++++++++---------------- 2 files changed, 31 insertions(+), 52 deletions(-) diff --git a/roborock/devices/traits/a01/__init__.py b/roborock/devices/traits/a01/__init__.py index 53103f54..1e3f44ae 100644 --- a/roborock/devices/traits/a01/__init__.py +++ b/roborock/devices/traits/a01/__init__.py @@ -92,19 +92,23 @@ } -def convert_dyad_value(protocol: int, value: Any) -> Any: +def convert_dyad_value(protocol_value: RoborockDyadDataProtocol, 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) + try: + return converter(value) + except (ValueError, TypeError): + return None return None -def convert_zeo_value(protocol: int, value: Any) -> Any: +def convert_zeo_value(protocol_value: RoborockZeoProtocol, 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) + try: + return converter(value) + except (ValueError, TypeError): + return None return None @@ -119,11 +123,7 @@ async def query_values(self, protocols: list[RoborockDyadDataProtocol]) -> dict[ """Query the device for the values of the given Dyad protocols.""" params = {RoborockDyadDataProtocol.ID_QUERY: str([int(p) for p in protocols])} response = await send_decoded_command(self._channel, params) - return { - RoborockDyadDataProtocol(int(k)): v - for k, val in response.items() - if (v := convert_dyad_value(int(k), val)) is not None - } + 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.""" @@ -144,11 +144,7 @@ async def query_values(self, protocols: list[RoborockZeoProtocol]) -> dict[Robor """Query the device for the values of the given protocols.""" params = {RoborockZeoProtocol.ID_QUERY: str([int(p) for p in protocols])} response = await send_decoded_command(self._channel, params) - return { - RoborockZeoProtocol(int(k)): v - for k, val in response.items() - if (v := convert_zeo_value(int(k), val)) is not None - } + 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.""" diff --git a/tests/devices/traits/a01/test_init.py b/tests/devices/traits/a01/test_init.py index 05601eff..a8d2a8c2 100644 --- a/tests/devices/traits/a01/test_init.py +++ b/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", }, ), ],