Skip to content

Commit 87d9aa6

Browse files
authored
feat: add v1 api support for the list of maps (#499)
* feat: add v1 api support for the list of maps * chore: add additional test coverage * chore: fix lint errors * fix: add a decorator to mark traits as mqtt only * chore: fix lint errors * chore: fix lint errors * chore: fix lint errors * chore: add comment describing the decorator check
1 parent 279283d commit 87d9aa6

File tree

10 files changed

+325
-14
lines changed

10 files changed

+325
-14
lines changed

roborock/cli.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -439,6 +439,16 @@ async def set_volume(ctx, device_id: str, volume: int):
439439
click.echo(f"Set Device {device_id} volume to {volume}")
440440

441441

442+
@session.command()
443+
@click.option("--device_id", required=True)
444+
@click.pass_context
445+
@async_command
446+
async def maps(ctx, device_id: str):
447+
"""Get device maps info."""
448+
context: RoborockContext = ctx.obj
449+
await _display_v1_trait(context, device_id, lambda v1: v1.maps)
450+
451+
442452
@click.command()
443453
@click.option("--device_id", required=True)
444454
@click.option("--cmd", required=True)
@@ -680,6 +690,7 @@ def write_markdown_table(product_features: dict[str, dict[str, any]], all_featur
680690
cli.add_command(clean_summary)
681691
cli.add_command(volume)
682692
cli.add_command(set_volume)
693+
cli.add_command(maps)
683694

684695

685696
def main():

roborock/devices/device_manager.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@ def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> Roborock
149149
match device.pv:
150150
case DeviceVersion.V1:
151151
channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
152-
trait = v1.create(product, channel.rpc_channel)
152+
trait = v1.create(product, channel.rpc_channel, channel.mqtt_rpc_channel)
153153
case DeviceVersion.A01:
154154
channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
155155
trait = a01.create(product, channel)

roborock/devices/traits/v1/__init__.py

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,29 @@
11
"""Create traits for V1 devices."""
22

3+
import logging
34
from dataclasses import dataclass, field, fields
45

5-
from roborock.containers import HomeDataProduct
6+
from roborock.containers import HomeData, HomeDataProduct
67
from roborock.devices.traits import Trait
78
from roborock.devices.v1_rpc_channel import V1RpcChannel
89

910
from .clean_summary import CleanSummaryTrait
1011
from .common import V1TraitMixin
1112
from .do_not_disturb import DoNotDisturbTrait
13+
from .maps import MapsTrait
1214
from .status import StatusTrait
1315
from .volume import SoundVolumeTrait
1416

17+
_LOGGER = logging.getLogger(__name__)
18+
1519
__all__ = [
1620
"create",
1721
"PropertiesApi",
1822
"StatusTrait",
1923
"DoNotDisturbTrait",
2024
"CleanSummaryTrait",
2125
"SoundVolumeTrait",
26+
"MapsTrait",
2227
]
2328

2429

@@ -34,12 +39,14 @@ class PropertiesApi(Trait):
3439
dnd: DoNotDisturbTrait
3540
clean_summary: CleanSummaryTrait
3641
sound_volume: SoundVolumeTrait
42+
maps: MapsTrait
3743

3844
# In the future optional fields can be added below based on supported features
3945

40-
def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel) -> None:
46+
def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel) -> None:
4147
"""Initialize the V1TraitProps with None values."""
4248
self.status = StatusTrait(product)
49+
self.maps = MapsTrait(self.status)
4350

4451
# This is a hack to allow setting the rpc_channel on all traits. This is
4552
# used so we can preserve the dataclass behavior when the values in the
@@ -49,9 +56,14 @@ def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel) -> None:
4956
if (trait := getattr(self, item.name, None)) is None:
5057
trait = item.type()
5158
setattr(self, item.name, trait)
52-
trait._rpc_channel = rpc_channel
59+
# The decorator `@common.mqtt_rpc_channel` means that the trait needs
60+
# to use the mqtt_rpc_channel (cloud only) instead of the rpc_channel (adaptive)
61+
if hasattr(trait, "mqtt_rpc_channel"):
62+
trait._rpc_channel = mqtt_rpc_channel
63+
else:
64+
trait._rpc_channel = rpc_channel
5365

5466

55-
def create(product: HomeDataProduct, rpc_channel: V1RpcChannel) -> PropertiesApi:
67+
def create(product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel) -> PropertiesApi:
5668
"""Create traits for V1 devices."""
57-
return PropertiesApi(product, rpc_channel)
69+
return PropertiesApi(product, rpc_channel, mqtt_rpc_channel)

roborock/devices/traits/v1/common.py

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,17 @@
33
This is an internal library and should not be used directly by consumers.
44
"""
55

6+
import logging
67
from abc import ABC
7-
from dataclasses import asdict, dataclass, fields
8+
from dataclasses import dataclass, fields
89
from typing import ClassVar, Self
910

1011
from roborock.containers import RoborockBase
1112
from roborock.devices.v1_rpc_channel import V1RpcChannel
1213
from roborock.roborock_typing import RoborockCommand
1314

15+
_LOGGER = logging.getLogger(__name__)
16+
1417
V1ResponseData = dict | list | int | str
1518

1619

@@ -77,9 +80,11 @@ async def refresh(self) -> Self:
7780
"""Refresh the contents of this trait."""
7881
response = await self.rpc_channel.send_command(self.command)
7982
new_data = self._parse_response(response)
80-
for k, v in asdict(new_data).items():
81-
if v is not None:
82-
setattr(self, k, v)
83+
if not isinstance(new_data, RoborockBase):
84+
raise ValueError(f"Internal error, unexpected response type: {new_data!r}")
85+
for field in fields(new_data):
86+
new_value = getattr(new_data, field.name, None)
87+
setattr(self, field.name, new_value)
8388
return self
8489

8590

@@ -113,3 +118,18 @@ def _parse_response(cls, response: V1ResponseData) -> Self:
113118
raise ValueError(f"Unexpected response format: {response!r}")
114119
value_field = _get_value_field(cls)
115120
return cls(**{value_field: response})
121+
122+
123+
def mqtt_rpc_channel(cls):
124+
"""Decorator to mark a function as cloud only.
125+
126+
Normally a trait uses an adaptive rpc channel that can use either local
127+
or cloud communication depending on what is available. This will force
128+
the trait to always use the cloud rpc channel.
129+
"""
130+
131+
def wrapper(*args, **kwargs):
132+
return cls(*args, **kwargs)
133+
134+
cls.mqtt_rpc_channel = True # type: ignore[attr-defined]
135+
return wrapper

roborock/devices/traits/v1/maps.py

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
"""Trait for managing maps and room mappings on Roborock devices.
2+
3+
New datatypes are introduced here to manage the additional information associated
4+
with maps and rooms, such as map names and room names. These override the
5+
base container datatypes to add additional fields.
6+
"""
7+
import logging
8+
from typing import Self
9+
10+
from roborock.containers import MultiMapsList, MultiMapsListMapInfo
11+
from roborock.devices.traits.v1 import common
12+
from roborock.roborock_typing import RoborockCommand
13+
14+
from .status import StatusTrait
15+
16+
_LOGGER = logging.getLogger(__name__)
17+
18+
19+
@common.mqtt_rpc_channel
20+
class MapsTrait(MultiMapsList, common.V1TraitMixin):
21+
"""Trait for managing the maps of Roborock devices.
22+
23+
A device may have multiple maps, each identified by a unique map_flag.
24+
Each map can have multiple rooms associated with it, in a `RoomMapping`.
25+
"""
26+
27+
command = RoborockCommand.GET_MULTI_MAPS_LIST
28+
29+
def __init__(self, status_trait: StatusTrait) -> None:
30+
"""Initialize the MapsTrait.
31+
32+
We keep track of the StatusTrait to ensure we have the latest
33+
status information when dealing with maps.
34+
"""
35+
super().__init__()
36+
self._status_trait = status_trait
37+
38+
@property
39+
def current_map(self) -> int | None:
40+
"""Returns the currently active map (map_flag), if available."""
41+
return self._status_trait.current_map
42+
43+
@property
44+
def current_map_info(self) -> MultiMapsListMapInfo | None:
45+
"""Returns the currently active map info, if available."""
46+
if (current_map := self.current_map) is None or self.map_info is None:
47+
return None
48+
for map_info in self.map_info:
49+
if map_info.map_flag == current_map:
50+
return map_info
51+
return None
52+
53+
async def set_current_map(self, map_flag: int) -> None:
54+
"""Update the current map of the device by it's map_flag id."""
55+
await self.rpc_channel.send_command(RoborockCommand.LOAD_MULTI_MAP, params=[map_flag])
56+
# Refresh our status to make sure it reflects the new map
57+
await self._status_trait.refresh()
58+
59+
def _parse_response(self, response: common.V1ResponseData) -> Self:
60+
"""Parse the response from the device into a MapsTrait instance.
61+
62+
This overrides the base implementation to handle the specific
63+
response format for the multi maps list. This is needed because we have
64+
a custom constructor that requires the StatusTrait.
65+
"""
66+
if not isinstance(response, list):
67+
raise ValueError(f"Unexpected MapsTrait response format: {response!r}")
68+
response = response[0]
69+
if not isinstance(response, dict):
70+
raise ValueError(f"Unexpected MapsTrait response format: {response!r}")
71+
return MultiMapsList.from_dict(response)

tests/devices/test_v1_device.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,13 +35,19 @@ def rpc_channel_fixture() -> AsyncMock:
3535
return AsyncMock()
3636

3737

38+
@pytest.fixture(autouse=True, name="mqtt_rpc_channel")
39+
def mqtt_rpc_channel_fixture() -> AsyncMock:
40+
"""Fixture to set up the channel for tests."""
41+
return AsyncMock()
42+
43+
3844
@pytest.fixture(autouse=True, name="device")
39-
def device_fixture(channel: AsyncMock, rpc_channel: AsyncMock) -> RoborockDevice:
45+
def device_fixture(channel: AsyncMock, rpc_channel: AsyncMock, mqtt_rpc_channel: AsyncMock) -> RoborockDevice:
4046
"""Fixture to set up the device for tests."""
4147
return RoborockDevice(
4248
device_info=HOME_DATA.devices[0],
4349
channel=channel,
44-
trait=v1.create(HOME_DATA.products[0], rpc_channel),
50+
trait=v1.create(HOME_DATA.products[0], rpc_channel, mqtt_rpc_channel),
4551
)
4652

4753

tests/devices/traits/v1/fixtures.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,17 @@ def rpc_channel_fixture() -> AsyncMock:
2727
return AsyncMock()
2828

2929

30+
@pytest.fixture(autouse=True, name="mock_mqtt_rpc_channel")
31+
def mqtt_rpc_channel_fixture() -> AsyncMock:
32+
"""Fixture to set up the channel for tests."""
33+
return AsyncMock()
34+
35+
3036
@pytest.fixture(autouse=True, name="device")
31-
def device_fixture(channel: AsyncMock, mock_rpc_channel: AsyncMock) -> RoborockDevice:
37+
def device_fixture(channel: AsyncMock, mock_rpc_channel: AsyncMock, mock_mqtt_rpc_channel: AsyncMock) -> RoborockDevice:
3238
"""Fixture to set up the device for tests."""
3339
return RoborockDevice(
3440
device_info=HOME_DATA.devices[0],
3541
channel=channel,
36-
trait=v1.create(HOME_DATA.products[0], mock_rpc_channel),
42+
trait=v1.create(HOME_DATA.products[0], mock_rpc_channel, mock_mqtt_rpc_channel),
3743
)

0 commit comments

Comments
 (0)