Skip to content

Commit f875e7a

Browse files
Lash-Lallenporter
andauthored
fix: convert a01 values (#647)
* fix: convert a01 values * chore: small changes to comments * chore: small changes to comments * fix: update tests and conversion logic --------- Co-authored-by: Allen Porter <allen.porter@gmail.com>
1 parent 580974c commit f875e7a

File tree

3 files changed

+179
-39
lines changed

3 files changed

+179
-39
lines changed

roborock/devices/traits/a01/__init__.py

Lines changed: 105 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,31 @@
1+
from collections.abc import Callable
2+
from datetime import time
13
from typing import Any
24

3-
from roborock.data import HomeDataProduct, RoborockCategory
5+
from roborock.data import DyadProductInfo, DyadSndState, HomeDataProduct, RoborockCategory
6+
from roborock.data.dyad.dyad_code_mappings import (
7+
DyadBrushSpeed,
8+
DyadCleanMode,
9+
DyadError,
10+
DyadSelfCleanLevel,
11+
DyadSelfCleanMode,
12+
DyadSuction,
13+
DyadWarmLevel,
14+
DyadWaterLevel,
15+
RoborockDyadStateCode,
16+
)
17+
from roborock.data.zeo.zeo_code_mappings import (
18+
ZeoDetergentType,
19+
ZeoDryingMode,
20+
ZeoError,
21+
ZeoMode,
22+
ZeoProgram,
23+
ZeoRinse,
24+
ZeoSoftenerType,
25+
ZeoSpin,
26+
ZeoState,
27+
ZeoTemperature,
28+
)
429
from roborock.devices.a01_channel import send_decoded_command
530
from roborock.devices.mqtt_channel import MqttChannel
631
from roborock.devices.traits import Trait
@@ -12,6 +37,81 @@
1237
]
1338

1439

40+
DYAD_PROTOCOL_ENTRIES: dict[RoborockDyadDataProtocol, Callable] = {
41+
RoborockDyadDataProtocol.STATUS: lambda val: RoborockDyadStateCode(val).name,
42+
RoborockDyadDataProtocol.SELF_CLEAN_MODE: lambda val: DyadSelfCleanMode(val).name,
43+
RoborockDyadDataProtocol.SELF_CLEAN_LEVEL: lambda val: DyadSelfCleanLevel(val).name,
44+
RoborockDyadDataProtocol.WARM_LEVEL: lambda val: DyadWarmLevel(val).name,
45+
RoborockDyadDataProtocol.CLEAN_MODE: lambda val: DyadCleanMode(val).name,
46+
RoborockDyadDataProtocol.SUCTION: lambda val: DyadSuction(val).name,
47+
RoborockDyadDataProtocol.WATER_LEVEL: lambda val: DyadWaterLevel(val).name,
48+
RoborockDyadDataProtocol.BRUSH_SPEED: lambda val: DyadBrushSpeed(val).name,
49+
RoborockDyadDataProtocol.POWER: lambda val: int(val),
50+
RoborockDyadDataProtocol.AUTO_DRY: lambda val: bool(val),
51+
RoborockDyadDataProtocol.MESH_LEFT: lambda val: int(360000 - val * 60),
52+
RoborockDyadDataProtocol.BRUSH_LEFT: lambda val: int(360000 - val * 60),
53+
RoborockDyadDataProtocol.ERROR: lambda val: DyadError(val).name,
54+
RoborockDyadDataProtocol.VOLUME_SET: lambda val: int(val),
55+
RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: lambda val: bool(val),
56+
RoborockDyadDataProtocol.AUTO_DRY_MODE: lambda val: bool(val),
57+
RoborockDyadDataProtocol.SILENT_DRY_DURATION: lambda val: int(val), # in minutes
58+
RoborockDyadDataProtocol.SILENT_MODE: lambda val: bool(val),
59+
RoborockDyadDataProtocol.SILENT_MODE_START_TIME: lambda val: time(
60+
hour=int(val / 60), minute=val % 60
61+
), # in minutes since 00:00
62+
RoborockDyadDataProtocol.SILENT_MODE_END_TIME: lambda val: time(
63+
hour=int(val / 60), minute=val % 60
64+
), # in minutes since 00:00
65+
RoborockDyadDataProtocol.RECENT_RUN_TIME: lambda val: [
66+
int(v) for v in val.split(",")
67+
], # minutes of cleaning in past few days.
68+
RoborockDyadDataProtocol.TOTAL_RUN_TIME: lambda val: int(val),
69+
RoborockDyadDataProtocol.SND_STATE: lambda val: DyadSndState.from_dict(val),
70+
RoborockDyadDataProtocol.PRODUCT_INFO: lambda val: DyadProductInfo.from_dict(val),
71+
}
72+
73+
ZEO_PROTOCOL_ENTRIES: dict[RoborockZeoProtocol, Callable] = {
74+
# read-only
75+
RoborockZeoProtocol.STATE: lambda val: ZeoState(val).name,
76+
RoborockZeoProtocol.COUNTDOWN: lambda val: int(val),
77+
RoborockZeoProtocol.WASHING_LEFT: lambda val: int(val),
78+
RoborockZeoProtocol.ERROR: lambda val: ZeoError(val).name,
79+
RoborockZeoProtocol.TIMES_AFTER_CLEAN: lambda val: int(val),
80+
RoborockZeoProtocol.DETERGENT_EMPTY: lambda val: bool(val),
81+
RoborockZeoProtocol.SOFTENER_EMPTY: lambda val: bool(val),
82+
# read-write
83+
RoborockZeoProtocol.MODE: lambda val: ZeoMode(val).name,
84+
RoborockZeoProtocol.PROGRAM: lambda val: ZeoProgram(val).name,
85+
RoborockZeoProtocol.TEMP: lambda val: ZeoTemperature(val).name,
86+
RoborockZeoProtocol.RINSE_TIMES: lambda val: ZeoRinse(val).name,
87+
RoborockZeoProtocol.SPIN_LEVEL: lambda val: ZeoSpin(val).name,
88+
RoborockZeoProtocol.DRYING_MODE: lambda val: ZeoDryingMode(val).name,
89+
RoborockZeoProtocol.DETERGENT_TYPE: lambda val: ZeoDetergentType(val).name,
90+
RoborockZeoProtocol.SOFTENER_TYPE: lambda val: ZeoSoftenerType(val).name,
91+
RoborockZeoProtocol.SOUND_SET: lambda val: bool(val),
92+
}
93+
94+
95+
def convert_dyad_value(protocol_value: RoborockDyadDataProtocol, value: Any) -> Any:
96+
"""Convert a dyad protocol value to its corresponding type."""
97+
if (converter := DYAD_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
98+
try:
99+
return converter(value)
100+
except (ValueError, TypeError):
101+
return None
102+
return None
103+
104+
105+
def convert_zeo_value(protocol_value: RoborockZeoProtocol, value: Any) -> Any:
106+
"""Convert a zeo protocol value to its corresponding type."""
107+
if (converter := ZEO_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
108+
try:
109+
return converter(value)
110+
except (ValueError, TypeError):
111+
return None
112+
return None
113+
114+
15115
class DyadApi(Trait):
16116
"""API for interacting with Dyad devices."""
17117

@@ -22,7 +122,8 @@ def __init__(self, channel: MqttChannel) -> None:
22122
async def query_values(self, protocols: list[RoborockDyadDataProtocol]) -> dict[RoborockDyadDataProtocol, Any]:
23123
"""Query the device for the values of the given Dyad protocols."""
24124
params = {RoborockDyadDataProtocol.ID_QUERY: str([int(p) for p in protocols])}
25-
return await send_decoded_command(self._channel, params)
125+
response = await send_decoded_command(self._channel, params)
126+
return {protocol: convert_dyad_value(protocol, response.get(protocol)) for protocol in protocols}
26127

27128
async def set_value(self, protocol: RoborockDyadDataProtocol, value: Any) -> dict[RoborockDyadDataProtocol, Any]:
28129
"""Set a value for a specific protocol on the device."""
@@ -42,7 +143,8 @@ def __init__(self, channel: MqttChannel) -> None:
42143
async def query_values(self, protocols: list[RoborockZeoProtocol]) -> dict[RoborockZeoProtocol, Any]:
43144
"""Query the device for the values of the given protocols."""
44145
params = {RoborockZeoProtocol.ID_QUERY: str([int(p) for p in protocols])}
45-
return await send_decoded_command(self._channel, params)
146+
response = await send_decoded_command(self._channel, params)
147+
return {protocol: convert_zeo_value(protocol, response.get(protocol)) for protocol in protocols}
46148

47149
async def set_value(self, protocol: RoborockZeoProtocol, value: Any) -> dict[RoborockZeoProtocol, Any]:
48150
"""Set a value for a specific protocol on the device."""

tests/devices/test_a01_traits.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
from unittest.mock import AsyncMock, Mock, patch
2+
3+
import pytest
4+
5+
from roborock.devices.traits.a01 import DyadApi, ZeoApi
6+
from roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol
7+
8+
9+
@pytest.fixture
10+
def mock_channel():
11+
channel = Mock()
12+
channel.send_command = AsyncMock()
13+
# Mocking send_decoded_command if it was a method on channel, but it's a standalone function imported in traits.
14+
# However, in traits/__init__.py it is imported as: from roborock.devices.a01_channel import send_decoded_command
15+
return channel
16+
17+
18+
@pytest.mark.asyncio
19+
async def test_dyad_query_values(mock_channel):
20+
with patch("roborock.devices.traits.a01.send_decoded_command", new_callable=AsyncMock) as mock_send:
21+
api = DyadApi(mock_channel)
22+
23+
# Setup mock return value (raw values)
24+
mock_send.return_value = {
25+
int(RoborockDyadDataProtocol.CLEAN_MODE): 1,
26+
int(RoborockDyadDataProtocol.POWER): 100,
27+
}
28+
29+
protocols = [RoborockDyadDataProtocol.CLEAN_MODE, RoborockDyadDataProtocol.POWER]
30+
result = await api.query_values(protocols)
31+
32+
# Verify conversion
33+
assert RoborockDyadDataProtocol.CLEAN_MODE in result
34+
assert RoborockDyadDataProtocol.POWER in result
35+
36+
assert isinstance(result[RoborockDyadDataProtocol.CLEAN_MODE], str)
37+
assert result[RoborockDyadDataProtocol.POWER] == 100
38+
39+
40+
@pytest.mark.asyncio
41+
async def test_zeo_query_values(mock_channel):
42+
with patch("roborock.devices.traits.a01.send_decoded_command", new_callable=AsyncMock) as mock_send:
43+
api = ZeoApi(mock_channel)
44+
45+
mock_send.return_value = {
46+
int(RoborockZeoProtocol.STATE): 6, # spinning
47+
int(RoborockZeoProtocol.COUNTDOWN): 120,
48+
}
49+
50+
protocols = [RoborockZeoProtocol.STATE, RoborockZeoProtocol.COUNTDOWN]
51+
result = await api.query_values(protocols)
52+
53+
assert RoborockZeoProtocol.STATE in result
54+
assert result[RoborockZeoProtocol.STATE] == "spinning"
55+
assert result[RoborockZeoProtocol.COUNTDOWN] == 120

tests/devices/traits/a01/test_init.py

Lines changed: 19 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import datetime
12
from collections.abc import Generator
23
from typing import Any
34
from unittest.mock import AsyncMock, call, patch
@@ -51,17 +52,16 @@ async def test_dyad_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMo
5152
]
5253
)
5354
assert result == {
54-
# Note: Bugs here, returning raw values
5555
RoborockDyadDataProtocol.POWER: 1,
56-
RoborockDyadDataProtocol.STATUS: 6,
57-
RoborockDyadDataProtocol.WATER_LEVEL: 3,
58-
RoborockDyadDataProtocol.MESH_LEFT: 120,
59-
RoborockDyadDataProtocol.BRUSH_LEFT: 90,
60-
RoborockDyadDataProtocol.SILENT_MODE_START_TIME: 85,
61-
RoborockDyadDataProtocol.RECENT_RUN_TIME: "3,4,5",
56+
RoborockDyadDataProtocol.STATUS: "self_clean_deep_cleaning",
57+
RoborockDyadDataProtocol.WATER_LEVEL: "l3",
58+
RoborockDyadDataProtocol.MESH_LEFT: 352800,
59+
RoborockDyadDataProtocol.BRUSH_LEFT: 354600,
60+
RoborockDyadDataProtocol.SILENT_MODE_START_TIME: datetime.time(1, 25),
61+
RoborockDyadDataProtocol.RECENT_RUN_TIME: [3, 4, 5],
6262
RoborockDyadDataProtocol.TOTAL_RUN_TIME: 123456,
63-
RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: 1,
64-
RoborockDyadDataProtocol.AUTO_DRY_MODE: 0,
63+
RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: True,
64+
RoborockDyadDataProtocol.AUTO_DRY_MODE: False,
6565
}
6666

6767
# 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
8686
9999: -3,
8787
},
8888
{
89-
# Note: Bug here, should return enum value
90-
RoborockDyadDataProtocol.STATUS: 3,
91-
# Note: Bug here, unknown value should not be returned
92-
7: 1,
93-
9999: -3,
89+
RoborockDyadDataProtocol.STATUS: "charging",
9490
},
9591
),
9692
(
@@ -99,8 +95,7 @@ async def test_dyad_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMo
9995
RoborockDyadDataProtocol.SILENT_MODE_START_TIME: "invalid",
10096
},
10197
{
102-
# Note: Bug here, invalid value should not be returned
103-
RoborockDyadDataProtocol.SILENT_MODE_START_TIME: "invalid",
98+
RoborockDyadDataProtocol.SILENT_MODE_START_TIME: None,
10499
},
105100
),
106101
(
@@ -111,11 +106,7 @@ async def test_dyad_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMo
111106
9999: -3,
112107
},
113108
{
114-
# Note: Bug here, should return time value
115-
RoborockDyadDataProtocol.SILENT_MODE_START_TIME: 85,
116-
# Note: Bug here, additional values should not be returned
117-
RoborockDyadDataProtocol.POWER: 2,
118-
9999: -3,
109+
RoborockDyadDataProtocol.SILENT_MODE_START_TIME: datetime.time(1, 25),
119110
},
120111
),
121112
],
@@ -164,10 +155,10 @@ async def test_zeo_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMoc
164155
)
165156
assert result == {
166157
# Note: Bug here, should return enum/bool values
167-
RoborockZeoProtocol.STATE: 6,
168-
RoborockZeoProtocol.TEMP: 3,
169-
RoborockZeoProtocol.DETERGENT_EMPTY: 1,
170-
RoborockZeoProtocol.SOFTENER_EMPTY: 0,
158+
RoborockZeoProtocol.STATE: "spinning",
159+
RoborockZeoProtocol.TEMP: "medium",
160+
RoborockZeoProtocol.DETERGENT_EMPTY: True,
161+
RoborockZeoProtocol.SOFTENER_EMPTY: False,
171162
RoborockZeoProtocol.TIMES_AFTER_CLEAN: 1,
172163
RoborockZeoProtocol.WASHING_LEFT: 0,
173164
}
@@ -193,11 +184,7 @@ async def test_zeo_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMoc
193184
9999: -3,
194185
},
195186
{
196-
# Note: Bug here, should return enum/bool value
197-
RoborockZeoProtocol.STATE: 1,
198-
# Note: Bug here, unknown value should not be returned
199-
7: 1,
200-
9999: -3,
187+
RoborockZeoProtocol.STATE: "standby",
201188
},
202189
),
203190
(
@@ -206,8 +193,7 @@ async def test_zeo_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMoc
206193
RoborockZeoProtocol.WASHING_LEFT: "invalid",
207194
},
208195
{
209-
# Note: Bug here, invalid value should not be returned
210-
RoborockZeoProtocol.WASHING_LEFT: "invalid",
196+
RoborockZeoProtocol.WASHING_LEFT: None,
211197
},
212198
),
213199
(
@@ -218,10 +204,7 @@ async def test_zeo_api_query_values(mock_channel: AsyncMock, mock_send: AsyncMoc
218204
9999: -3,
219205
},
220206
{
221-
RoborockZeoProtocol.STATE: 1,
222-
# Note: Bug here, these values were not requested and should not be returned
223-
RoborockZeoProtocol.WASHING_LEFT: 2,
224-
9999: -3,
207+
RoborockZeoProtocol.STATE: "standby",
225208
},
226209
),
227210
],

0 commit comments

Comments
 (0)