Skip to content

Commit 1585e1c

Browse files
authored
feat: add a clean summary trait (#476)
1 parent 5a2ca23 commit 1585e1c

File tree

11 files changed

+323
-71
lines changed

11 files changed

+323
-71
lines changed

roborock/cli.py

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -388,16 +388,32 @@ async def status(ctx, device_id: str):
388388
device_manager = await context.get_device_manager()
389389
device = await device_manager.get_device(device_id)
390390

391-
click.echo(f"Getting status for device {device_id}")
392391
if not (status_trait := device.traits.get("status")):
393392
click.echo(f"Device {device.name} does not have a status trait")
394393
return
395394

396395
status_result = await status_trait.get_status()
397-
click.echo(f"Device {device_id} status:")
398396
click.echo(dump_json(status_result.as_dict()))
399397

400398

399+
@session.command()
400+
@click.option("--device_id", required=True)
401+
@click.pass_context
402+
@async_command
403+
async def clean_summary(ctx, device_id: str):
404+
"""Get device clean summary."""
405+
context: RoborockContext = ctx.obj
406+
407+
device_manager = await context.get_device_manager()
408+
device = await device_manager.get_device(device_id)
409+
if not (clean_summary_trait := device.traits.get("clean_summary")):
410+
click.echo(f"Device {device.name} does not have a clean summary trait")
411+
return
412+
413+
clean_summary_result = await clean_summary_trait.get_clean_summary()
414+
click.echo(dump_json(clean_summary_result.as_dict()))
415+
416+
401417
@click.command()
402418
@click.option("--device_id", required=True)
403419
@click.option("--cmd", required=True)
@@ -636,6 +652,7 @@ def write_markdown_table(product_features: dict[str, dict[str, any]], all_featur
636652
cli.add_command(session)
637653
cli.add_command(get_device_info)
638654
cli.add_command(update_docs)
655+
cli.add_command(clean_summary)
639656

640657

641658
def main():

roborock/devices/device_manager.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from .channel import Channel
2323
from .mqtt_channel import create_mqtt_channel
2424
from .traits.b01.props import B01PropsApi
25+
from .traits.clean_summary import CleanSummaryTrait
2526
from .traits.dnd import DoNotDisturbTrait
2627
from .traits.dyad import DyadApi
2728
from .traits.status import StatusTrait
@@ -154,6 +155,7 @@ def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> Roborock
154155
channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
155156
traits.append(StatusTrait(product, channel.rpc_channel))
156157
traits.append(DoNotDisturbTrait(channel.rpc_channel))
158+
traits.append(CleanSummaryTrait(channel.rpc_channel))
157159
case DeviceVersion.A01:
158160
mqtt_channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
159161
match product.category:
Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
"""Module for Roborock V1 devices.
2+
3+
This interface is experimental and subject to breaking changes without notice
4+
until the API is stable.
5+
"""
6+
7+
import logging
8+
9+
from roborock.containers import (
10+
CleanSummary,
11+
)
12+
from roborock.devices.v1_rpc_channel import V1RpcChannel
13+
from roborock.roborock_typing import RoborockCommand
14+
from roborock.util import unpack_list
15+
16+
from .trait import Trait
17+
18+
_LOGGER = logging.getLogger(__name__)
19+
20+
__all__ = [
21+
"CleanSummaryTrait",
22+
]
23+
24+
25+
class CleanSummaryTrait(Trait):
26+
"""Trait for managing the clean summary of Roborock devices."""
27+
28+
name = "clean_summary"
29+
30+
def __init__(self, rpc_channel: V1RpcChannel) -> None:
31+
"""Initialize the CleanSummaryTrait."""
32+
self._rpc_channel = rpc_channel
33+
34+
async def get_clean_summary(self) -> CleanSummary:
35+
"""Get the current clean summary of the device.
36+
37+
This is a placeholder command and will likely be changed/moved in the future.
38+
"""
39+
clean_summary = await self._rpc_channel.send_command(RoborockCommand.GET_CLEAN_SUMMARY)
40+
if isinstance(clean_summary, dict):
41+
return CleanSummary.from_dict(clean_summary)
42+
elif isinstance(clean_summary, list):
43+
clean_time, clean_area, clean_count, records = unpack_list(clean_summary, 4)
44+
return CleanSummary(
45+
clean_time=clean_time,
46+
clean_area=clean_area,
47+
clean_count=clean_count,
48+
records=records,
49+
)
50+
elif isinstance(clean_summary, int):
51+
return CleanSummary(clean_time=clean_summary)
52+
raise ValueError(f"Unexpected clean summary format: {clean_summary!r}")

roborock/devices/traits/status.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
"""
66

77
import logging
8+
from typing import Any
89

910
from roborock.containers import (
1011
HomeDataProduct,
@@ -40,4 +41,9 @@ async def get_status(self) -> Status:
4041
This is a placeholder command and will likely be changed/moved in the future.
4142
"""
4243
status_type: type[Status] = ModelStatus.get(self._product_info.model, S7MaxVStatus)
43-
return await self._rpc_channel.send_command(RoborockCommand.GET_STATUS, response_type=status_type)
44+
status: dict[str, Any] | list = await self._rpc_channel.send_command(RoborockCommand.GET_STATUS)
45+
if isinstance(status, list):
46+
status = status[0]
47+
if not isinstance(status, dict):
48+
raise ValueError(f"Unexpected status format: {status!r}")
49+
return status_type.from_dict(status)

roborock/devices/v1_rpc_channel.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
CommandType,
1818
ParamsType,
1919
RequestMessage,
20+
ResponseData,
2021
SecurityData,
2122
decode_rpc_response,
2223
)
@@ -130,15 +131,15 @@ async def _send_raw_command(
130131
method: CommandType,
131132
*,
132133
params: ParamsType = None,
133-
) -> Any:
134+
) -> ResponseData:
134135
"""Send a command and return a parsed response RoborockBase type."""
135136
request_message = RequestMessage(method, params=params)
136137
_LOGGER.debug(
137138
"Sending command (%s, request_id=%s): %s, params=%s", self._name, request_message.request_id, method, params
138139
)
139140
message = self._payload_encoder(request_message)
140141

141-
future: asyncio.Future[dict[str, Any]] = asyncio.Future()
142+
future: asyncio.Future[ResponseData] = asyncio.Future()
142143

143144
def find_response(response_message: RoborockMessage) -> None:
144145
try:

roborock/protocols/v1_protocol.py

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -95,15 +95,18 @@ def _as_payload(self, security_data: SecurityData | None) -> bytes:
9595
)
9696

9797

98+
ResponseData = dict[str, Any] | list | int
99+
100+
98101
@dataclass(kw_only=True, frozen=True)
99102
class ResponseMessage:
100103
"""Data structure for v1 RoborockMessage responses."""
101104

102105
request_id: int | None
103106
"""The request ID of the response."""
104107

105-
data: dict[str, Any]
106-
"""The data of the response."""
108+
data: ResponseData
109+
"""The data of the response, where the type depends on the command."""
107110

108111

109112
def decode_rpc_response(message: RoborockMessage) -> ResponseMessage:
@@ -139,12 +142,12 @@ def decode_rpc_response(message: RoborockMessage) -> ResponseMessage:
139142
if not (result := data_point_response.get("result")):
140143
raise RoborockException(f"Invalid V1 message format: missing 'result' in data point for {message.payload!r}")
141144
_LOGGER.debug("Decoded V1 message result: %s", result)
142-
if isinstance(result, list) and result:
143-
result = result[0]
144145
if isinstance(result, str) and result == "ok":
145146
result = {}
146-
if not isinstance(result, dict):
147-
raise RoborockException(f"Invalid V1 message format: 'result' should be a dictionary for {message.payload!r}")
147+
if not isinstance(result, (dict, list, int)):
148+
raise RoborockException(
149+
f"Invalid V1 message format: 'result' was unexpected type {type(result)}. {message.payload!r}"
150+
)
148151
return ResponseMessage(request_id=request_id, data=result)
149152

150153

tests/devices/test_v1_device.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ async def test_device_connection(device: RoborockDevice, channel: AsyncMock) ->
7373
async def test_device_get_status_command(device: RoborockDevice, rpc_channel: AsyncMock) -> None:
7474
"""Test the device get_status command."""
7575
# Mock response for get_status command
76-
rpc_channel.send_command.return_value = STATUS
76+
rpc_channel.send_command.return_value = [STATUS.as_dict()]
7777

7878
# Test get_status and verify the command was sent
7979
status_api = device.traits["status"]
Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
"""Tests for the CleanSummary class."""
2+
3+
from unittest.mock import AsyncMock
4+
5+
import pytest
6+
7+
from roborock.containers import CleanSummary
8+
from roborock.devices.traits.clean_summary import CleanSummaryTrait
9+
from roborock.devices.v1_rpc_channel import V1RpcChannel
10+
from roborock.exceptions import RoborockException
11+
from roborock.roborock_typing import RoborockCommand
12+
13+
CLEAN_SUMMARY_DATA = [
14+
1442559,
15+
24258125000,
16+
296,
17+
[
18+
1756848207,
19+
1754930385,
20+
1753203976,
21+
1752183435,
22+
1747427370,
23+
1746204046,
24+
1745601543,
25+
1744387080,
26+
1743528522,
27+
1742489154,
28+
1741022299,
29+
1740433682,
30+
1739902516,
31+
1738875106,
32+
1738864366,
33+
1738620067,
34+
1736873889,
35+
1736197544,
36+
1736121269,
37+
1734458038,
38+
],
39+
]
40+
41+
42+
@pytest.fixture
43+
def mock_rpc_channel() -> AsyncMock:
44+
"""Create a mock RPC channel."""
45+
mock_channel = AsyncMock(spec=V1RpcChannel)
46+
# Ensure send_command is an AsyncMock that returns awaitable coroutines
47+
mock_channel.send_command = AsyncMock()
48+
return mock_channel
49+
50+
51+
@pytest.fixture
52+
def clean_summary_trait(mock_rpc_channel: AsyncMock) -> CleanSummaryTrait:
53+
"""Create a CleanSummaryTrait instance with mocked dependencies."""
54+
return CleanSummaryTrait(mock_rpc_channel)
55+
56+
57+
@pytest.fixture
58+
def sample_clean_summary() -> CleanSummary:
59+
"""Create a sample CleanSummary for testing."""
60+
return CleanSummary(
61+
clean_area=100,
62+
clean_time=3600,
63+
)
64+
65+
66+
def test_trait_name(clean_summary_trait: CleanSummaryTrait) -> None:
67+
"""Test that the trait has the correct name."""
68+
assert clean_summary_trait.name == "clean_summary"
69+
70+
71+
async def test_get_clean_summary_success(
72+
clean_summary_trait: CleanSummaryTrait, mock_rpc_channel: AsyncMock, sample_clean_summary: CleanSummary
73+
) -> None:
74+
"""Test successfully getting clean summary."""
75+
# Setup mock to return the sample clean summary
76+
mock_rpc_channel.send_command.return_value = CLEAN_SUMMARY_DATA
77+
78+
# Call the method
79+
result = await clean_summary_trait.get_clean_summary()
80+
81+
# Verify the result
82+
assert result.clean_area == 24258125000
83+
assert result.clean_time == 1442559
84+
assert result.square_meter_clean_area == 24258.1
85+
assert result.clean_count == 296
86+
assert result.records
87+
assert len(result.records) == 20
88+
assert result.records[0] == 1756848207
89+
90+
# Verify the RPC call was made correctly
91+
mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.GET_CLEAN_SUMMARY)
92+
93+
94+
async def test_get_clean_summary_clean_time_only(
95+
clean_summary_trait: CleanSummaryTrait, mock_rpc_channel: AsyncMock, sample_clean_summary: CleanSummary
96+
) -> None:
97+
"""Test successfully getting clean summary where the response only has the clean time."""
98+
99+
mock_rpc_channel.send_command.return_value = [1442559]
100+
101+
# Call the method
102+
result = await clean_summary_trait.get_clean_summary()
103+
104+
# Verify the result
105+
assert result.clean_area is None
106+
assert result.clean_time == 1442559
107+
assert result.square_meter_clean_area is None
108+
assert result.clean_count is None
109+
assert not result.records
110+
111+
# Verify the RPC call was made correctly
112+
mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.GET_CLEAN_SUMMARY)
113+
114+
115+
async def test_get_clean_summary_propagates_exception(
116+
clean_summary_trait: CleanSummaryTrait, mock_rpc_channel: AsyncMock
117+
) -> None:
118+
"""Test that exceptions from RPC channel are propagated in get_clean_summary."""
119+
120+
# Setup mock to raise an exception
121+
mock_rpc_channel.send_command.side_effect = RoborockException("Communication error")
122+
123+
# Verify the exception is propagated
124+
with pytest.raises(RoborockException, match="Communication error"):
125+
await clean_summary_trait.get_clean_summary()

0 commit comments

Comments
 (0)