diff --git a/roborock/devices/traits/a01/__init__.py b/roborock/devices/traits/a01/__init__.py index e804f23c..1e3f44ae 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,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.""" diff --git a/tests/devices/test_a01_traits.py b/tests/devices/test_a01_traits.py new file mode 100644 index 00000000..be400f0a --- /dev/null +++ b/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 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", }, ), ],