From 0fbeca34e644e9d497624e46b31ab3dc6ed97018 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sat, 27 Sep 2025 15:27:31 -0700 Subject: [PATCH 1/8] feat: add v1 api support for the list of maps --- roborock/cli.py | 11 ++ roborock/devices/traits/v1/__init__.py | 6 +- roborock/devices/traits/v1/common.py | 10 +- roborock/devices/traits/v1/maps.py | 70 +++++++++++ tests/devices/traits/v1/test_maps.py | 117 ++++++++++++++++++ .../__snapshots__/test_v1_protocol.ambr | 28 +++++ .../v1_protocol/get_multi_maps_list.json | 1 + 7 files changed, 238 insertions(+), 5 deletions(-) create mode 100644 roborock/devices/traits/v1/maps.py create mode 100644 tests/devices/traits/v1/test_maps.py create mode 100644 tests/protocols/testdata/v1_protocol/get_multi_maps_list.json diff --git a/roborock/cli.py b/roborock/cli.py index c5e39d44..acc17e69 100644 --- a/roborock/cli.py +++ b/roborock/cli.py @@ -439,6 +439,16 @@ async def set_volume(ctx, device_id: str, volume: int): click.echo(f"Set Device {device_id} volume to {volume}") +@session.command() +@click.option("--device_id", required=True) +@click.pass_context +@async_command +async def maps(ctx, device_id: str): + """Get device maps info.""" + context: RoborockContext = ctx.obj + await _display_v1_trait(context, device_id, lambda v1: v1.maps) + + @click.command() @click.option("--device_id", required=True) @click.option("--cmd", required=True) @@ -680,6 +690,7 @@ def write_markdown_table(product_features: dict[str, dict[str, any]], all_featur cli.add_command(clean_summary) cli.add_command(volume) cli.add_command(set_volume) +cli.add_command(maps) def main(): diff --git a/roborock/devices/traits/v1/__init__.py b/roborock/devices/traits/v1/__init__.py index e9e37eb4..9b3af76d 100644 --- a/roborock/devices/traits/v1/__init__.py +++ b/roborock/devices/traits/v1/__init__.py @@ -2,13 +2,14 @@ from dataclasses import dataclass, field, fields -from roborock.containers import HomeDataProduct +from roborock.containers import HomeData, HomeDataProduct from roborock.devices.traits import Trait from roborock.devices.v1_rpc_channel import V1RpcChannel from .clean_summary import CleanSummaryTrait from .common import V1TraitMixin from .do_not_disturb import DoNotDisturbTrait +from .maps import MapsTrait from .status import StatusTrait from .volume import SoundVolumeTrait @@ -19,6 +20,7 @@ "DoNotDisturbTrait", "CleanSummaryTrait", "SoundVolumeTrait", + "MapsTrait", ] @@ -34,12 +36,14 @@ class PropertiesApi(Trait): dnd: DoNotDisturbTrait clean_summary: CleanSummaryTrait sound_volume: SoundVolumeTrait + maps: MapsTrait # In the future optional fields can be added below based on supported features def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel) -> None: """Initialize the V1TraitProps with None values.""" self.status = StatusTrait(product) + self.maps = MapsTrait(self.status) # This is a hack to allow setting the rpc_channel on all traits. This is # used so we can preserve the dataclass behavior when the values in the diff --git a/roborock/devices/traits/v1/common.py b/roborock/devices/traits/v1/common.py index ad724947..7b14d0d3 100644 --- a/roborock/devices/traits/v1/common.py +++ b/roborock/devices/traits/v1/common.py @@ -4,7 +4,7 @@ """ from abc import ABC -from dataclasses import asdict, dataclass, fields +from dataclasses import dataclass, fields from typing import ClassVar, Self from roborock.containers import RoborockBase @@ -77,9 +77,11 @@ async def refresh(self) -> Self: """Refresh the contents of this trait.""" response = await self.rpc_channel.send_command(self.command) new_data = self._parse_response(response) - for k, v in asdict(new_data).items(): - if v is not None: - setattr(self, k, v) + if not isinstance(new_data, RoborockBase): + raise ValueError(f"Internal error, unexpected response type: {new_data!r}") + for field in fields(new_data): + new_value = getattr(new_data, field.name, None) + setattr(self, field.name, new_value) return self diff --git a/roborock/devices/traits/v1/maps.py b/roborock/devices/traits/v1/maps.py new file mode 100644 index 00000000..2da32de0 --- /dev/null +++ b/roborock/devices/traits/v1/maps.py @@ -0,0 +1,70 @@ +"""Trait for managing maps and room mappings on Roborock devices. + +New datatypes are introduced here to manage the additional information associated +with maps and rooms, such as map names and room names. These override the +base container datatypes to add additional fields. +""" +import logging +from typing import Self + +from roborock.containers import MultiMapsList, MultiMapsListMapInfo +from roborock.devices.traits.v1 import common +from roborock.roborock_typing import RoborockCommand + +from .status import StatusTrait + +_LOGGER = logging.getLogger(__name__) + + +class MapsTrait(MultiMapsList, common.V1TraitMixin): + """Trait for managing the maps of Roborock devices. + + A device may have multiple maps, each identified by a unique map_flag. + Each map can have multiple rooms associated with it, in a `RoomMapping`. + """ + + command = RoborockCommand.GET_MULTI_MAPS_LIST + + def __init__(self, status_trait: StatusTrait) -> None: + """Initialize the MapsTrait. + + We keep track of the StatusTrait to ensure we have the latest + status information when dealing with maps. + """ + super().__init__() + self._status_trait = status_trait + + @property + def current_map(self) -> int | None: + """Returns the currently active map (map_flag), if available.""" + return self._status_trait.current_map + + @property + def current_map_info(self) -> MultiMapsListMapInfo | None: + """Returns the currently active map info, if available.""" + if (current_map := self.current_map) is None or self.map_info is None: + return None + for map_info in self.map_info: + if map_info.map_flag == current_map: + return map_info + return None + + async def set_current_map(self, map_flag: int) -> None: + """Update the current map of the device by it's map_flag id.""" + await self.rpc_channel.send_command(RoborockCommand.LOAD_MULTI_MAP, params=[map_flag]) + # Refresh our status to make sure it reflects the new map + await self._status_trait.refresh() + + def _parse_response(self, response: common.V1ResponseData) -> Self: + """Parse the response from the device into a MapsTrait instance. + + This overrides the base implementation to handle the specific + response format for the multi maps list. This is needed because we have + a custom constructor that requires the StatusTrait. + """ + if not isinstance(response, list): + raise ValueError(f"Unexpected MapsTrait response format: {response!r}") + response = response[0] + if not isinstance(response, dict): + raise ValueError(f"Unexpected MapsTrait response format: {response!r}") + return MultiMapsList.from_dict(response) diff --git a/tests/devices/traits/v1/test_maps.py b/tests/devices/traits/v1/test_maps.py new file mode 100644 index 00000000..92881f10 --- /dev/null +++ b/tests/devices/traits/v1/test_maps.py @@ -0,0 +1,117 @@ +"""Tests for the Maps related functionality.""" + +from unittest.mock import AsyncMock + +import pytest + +from roborock.devices.device import RoborockDevice +from roborock.devices.traits.v1.maps import MapsTrait +from roborock.devices.traits.v1.status import StatusTrait +from roborock.roborock_typing import RoborockCommand +from tests import mock_data + +UPDATED_STATUS = { + **mock_data.STATUS, + "map_status": 123 * 4 + 3, # Set current map to 123 +} + + +MULTI_MAP_LIST_DATA = [ + { + "max_multi_map": 1, + "max_bak_map": 1, + "multi_map_count": 1, + "map_info": [ + { + "mapFlag": 0, + "add_time": 1747132930, + "length": 0, + "name": "Map 1", + "bak_maps": [{"mapFlag": 4, "add_time": 1747132936}], + }, + { + "mapFlag": 123, + "add_time": 1747132930, + "length": 0, + "name": "Map 1", + "bak_maps": [{"mapFlag": 4, "add_time": 1747132936}], + }, + ], + } +] + + +@pytest.fixture +def status_trait(device: RoborockDevice) -> StatusTrait: + """Create a MapsTrait instance with mocked dependencies.""" + assert device.v1_properties + return device.v1_properties.status + + +@pytest.fixture +def maps_trait(device: RoborockDevice) -> MapsTrait: + """Create a MapsTrait instance with mocked dependencies.""" + assert device.v1_properties + return device.v1_properties.maps + + +async def test_refresh_maps_trait( + maps_trait: MapsTrait, + mock_rpc_channel: AsyncMock, +) -> None: + """Test successfully getting multi maps list.""" + # Setup mock to return the sample multi maps list + mock_rpc_channel.send_command.return_value = MULTI_MAP_LIST_DATA + + # Call the method + await maps_trait.refresh() + + assert maps_trait.max_multi_map == 1 + assert maps_trait.max_bak_map == 1 + assert maps_trait.multi_map_count == 1 + assert maps_trait.map_info + assert len(maps_trait.map_info) == 2 + map_infos = maps_trait.map_info + assert len(map_infos) == 2 + assert map_infos[0].map_flag == 0 + assert map_infos[0].name == "Map 1" + assert map_infos[0].add_time == 1747132930 + assert map_infos[1].map_flag == 123 + assert map_infos[1].name == "Map 1" + assert map_infos[1].add_time == 1747132930 + + # Verify the RPC call was made correctly + mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.GET_MULTI_MAPS_LIST) + + +async def test_set_current_map( + status_trait: StatusTrait, + maps_trait: MapsTrait, + mock_rpc_channel: AsyncMock, +) -> None: + """Test successfully setting the current map.""" + mock_rpc_channel.send_command.side_effect = [ + mock_data.STATUS, # Initial status fetch + MULTI_MAP_LIST_DATA, # Response for LOAD_MULTI_MAP + {}, # Response for setting the current map + UPDATED_STATUS, # Response for refreshing status + ] + await status_trait.refresh() + + # First refresh to populate initial state + await maps_trait.refresh() + + # Call the method to set current map + await maps_trait.set_current_map(123) + + # Verify the current map is updated + assert maps_trait.current_map == 123 + + # Verify the RPC call was made correctly to load the map + mock_rpc_channel.send_command.assert_any_call(RoborockCommand.LOAD_MULTI_MAP, params=[123]) + # Command sent are: + # 1. GET_STATUS to get initial status + # 2. GET_MULTI_MAPS_LIST to get the map list + # 3. LOAD_MULTI_MAP to set the map + # 4. GET_STATUS to refresh the current map in status + assert mock_rpc_channel.send_command.call_count == 4 diff --git a/tests/protocols/__snapshots__/test_v1_protocol.ambr b/tests/protocols/__snapshots__/test_v1_protocol.ambr index df4ed552..cc31e868 100644 --- a/tests/protocols/__snapshots__/test_v1_protocol.ambr +++ b/tests/protocols/__snapshots__/test_v1_protocol.ambr @@ -49,6 +49,34 @@ ] ''' # --- +# name: test_decode_rpc_payload[get_multi_maps_list] + 20001 +# --- +# name: test_decode_rpc_payload[get_multi_maps_list].1 + ''' + [ + { + "max_multi_map": 1, + "max_bak_map": 1, + "multi_map_count": 1, + "map_info": [ + { + "mapFlag": 0, + "add_time": 1747132930, + "length": 0, + "name": "", + "bak_maps": [ + { + "mapFlag": 4, + "add_time": 1747132936 + } + ] + } + ] + } + ] + ''' +# --- # name: test_decode_rpc_payload[get_status] 20001 # --- diff --git a/tests/protocols/testdata/v1_protocol/get_multi_maps_list.json b/tests/protocols/testdata/v1_protocol/get_multi_maps_list.json new file mode 100644 index 00000000..f1e908a7 --- /dev/null +++ b/tests/protocols/testdata/v1_protocol/get_multi_maps_list.json @@ -0,0 +1 @@ +{"t":1758987228,"dps":{"102":"{\"id\":20001,\"result\":[{\"max_multi_map\":1,\"max_bak_map\":1,\"multi_map_count\":1,\"map_info\":[{\"mapFlag\":0,\"add_time\":1747132930,\"length\":0,\"name\":\"\",\"bak_maps\":[{\"mapFlag\":4,\"add_time\":1747132936}]}]}]}"}} From 432bbb6e61eecd95ece2a558fe879994ba1abad5 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sat, 27 Sep 2025 16:40:19 -0700 Subject: [PATCH 2/8] chore: add additional test coverage --- tests/devices/traits/v1/test_maps.py | 46 ++++++++++++++++++++++------ 1 file changed, 37 insertions(+), 9 deletions(-) diff --git a/tests/devices/traits/v1/test_maps.py b/tests/devices/traits/v1/test_maps.py index 92881f10..553e1e92 100644 --- a/tests/devices/traits/v1/test_maps.py +++ b/tests/devices/traits/v1/test_maps.py @@ -33,7 +33,7 @@ "mapFlag": 123, "add_time": 1747132930, "length": 0, - "name": "Map 1", + "name": "Map 2", "bak_maps": [{"mapFlag": 4, "add_time": 1747132936}], }, ], @@ -58,18 +58,28 @@ def maps_trait(device: RoborockDevice) -> MapsTrait: async def test_refresh_maps_trait( maps_trait: MapsTrait, mock_rpc_channel: AsyncMock, + status_trait: StatusTrait, ) -> None: """Test successfully getting multi maps list.""" # Setup mock to return the sample multi maps list - mock_rpc_channel.send_command.return_value = MULTI_MAP_LIST_DATA + mock_rpc_channel.send_command.side_effect = [ + mock_data.STATUS, # Initial status fetch + MULTI_MAP_LIST_DATA + ] + await status_trait.refresh() + # Populating the status information gives us the current map + # flag, but we have not loaded the rest of the information. + assert maps_trait.current_map == 0 + assert maps_trait.current_map_info is None - # Call the method + # Load the maps information await maps_trait.refresh() assert maps_trait.max_multi_map == 1 assert maps_trait.max_bak_map == 1 assert maps_trait.multi_map_count == 1 assert maps_trait.map_info + assert len(maps_trait.map_info) == 2 map_infos = maps_trait.map_info assert len(map_infos) == 2 @@ -77,12 +87,19 @@ async def test_refresh_maps_trait( assert map_infos[0].name == "Map 1" assert map_infos[0].add_time == 1747132930 assert map_infos[1].map_flag == 123 - assert map_infos[1].name == "Map 1" + assert map_infos[1].name == "Map 2" assert map_infos[1].add_time == 1747132930 - # Verify the RPC call was made correctly - mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.GET_MULTI_MAPS_LIST) + assert maps_trait.current_map == 0 + assert maps_trait.current_map_info is not None + assert maps_trait.current_map_info.map_flag == 0 + assert maps_trait.current_map_info.name == "Map 1" + # Verify the RPC call was made correctly + assert mock_rpc_channel.send_command.call_count == 2 + mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_STATUS) + mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_MULTI_MAPS_LIST) + async def test_set_current_map( status_trait: StatusTrait, @@ -101,17 +118,28 @@ async def test_set_current_map( # First refresh to populate initial state await maps_trait.refresh() + # Verify current map + + assert maps_trait.current_map == 0 + assert maps_trait.current_map_info + assert maps_trait.current_map_info.map_flag == 0 + assert maps_trait.current_map_info.name == "Map 1" + # Call the method to set current map await maps_trait.set_current_map(123) # Verify the current map is updated assert maps_trait.current_map == 123 + assert maps_trait.current_map_info + assert maps_trait.current_map_info.map_flag == 123 + assert maps_trait.current_map_info.name == "Map 2" - # Verify the RPC call was made correctly to load the map - mock_rpc_channel.send_command.assert_any_call(RoborockCommand.LOAD_MULTI_MAP, params=[123]) - # Command sent are: + # Verify the command sent are: # 1. GET_STATUS to get initial status # 2. GET_MULTI_MAPS_LIST to get the map list # 3. LOAD_MULTI_MAP to set the map # 4. GET_STATUS to refresh the current map in status assert mock_rpc_channel.send_command.call_count == 4 + mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_STATUS) + mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_MULTI_MAPS_LIST) + mock_rpc_channel.send_command.assert_any_call(RoborockCommand.LOAD_MULTI_MAP, params=[123]) From 8f25a60895e0c94f5e154adb4c5babfc192420b7 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sat, 27 Sep 2025 17:52:05 -0700 Subject: [PATCH 3/8] chore: fix lint errors --- tests/devices/traits/v1/test_maps.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/devices/traits/v1/test_maps.py b/tests/devices/traits/v1/test_maps.py index 553e1e92..02caef0d 100644 --- a/tests/devices/traits/v1/test_maps.py +++ b/tests/devices/traits/v1/test_maps.py @@ -64,7 +64,7 @@ async def test_refresh_maps_trait( # Setup mock to return the sample multi maps list mock_rpc_channel.send_command.side_effect = [ mock_data.STATUS, # Initial status fetch - MULTI_MAP_LIST_DATA + MULTI_MAP_LIST_DATA, ] await status_trait.refresh() # Populating the status information gives us the current map @@ -99,7 +99,7 @@ async def test_refresh_maps_trait( assert mock_rpc_channel.send_command.call_count == 2 mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_STATUS) mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_MULTI_MAPS_LIST) - + async def test_set_current_map( status_trait: StatusTrait, From 878607918dfb1a48bee1b7ce0ec38da578a053c6 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sat, 27 Sep 2025 21:03:39 -0700 Subject: [PATCH 4/8] fix: add a decorator to mark traits as mqtt only --- roborock/devices/device_manager.py | 2 +- roborock/devices/traits/v1/__init__.py | 16 ++++++++++++---- roborock/devices/traits/v1/common.py | 17 +++++++++++++++++ roborock/devices/traits/v1/maps.py | 1 + tests/devices/test_v1_device.py | 9 +++++++-- tests/devices/traits/v1/fixtures.py | 10 ++++++++-- tests/devices/traits/v1/test_maps.py | 23 +++++++++++++++++------ 7 files changed, 63 insertions(+), 15 deletions(-) diff --git a/roborock/devices/device_manager.py b/roborock/devices/device_manager.py index 29756f5e..66dfa25c 100644 --- a/roborock/devices/device_manager.py +++ b/roborock/devices/device_manager.py @@ -149,7 +149,7 @@ def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> Roborock match device.pv: case DeviceVersion.V1: channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache) - trait = v1.create(product, channel.rpc_channel) + trait = v1.create(product, channel.rpc_channel, channel.mqtt_rpc_channel) case DeviceVersion.A01: channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device) trait = a01.create(product, channel) diff --git a/roborock/devices/traits/v1/__init__.py b/roborock/devices/traits/v1/__init__.py index 9b3af76d..f1811af5 100644 --- a/roborock/devices/traits/v1/__init__.py +++ b/roborock/devices/traits/v1/__init__.py @@ -1,6 +1,7 @@ """Create traits for V1 devices.""" from dataclasses import dataclass, field, fields +import logging from roborock.containers import HomeData, HomeDataProduct from roborock.devices.traits import Trait @@ -13,6 +14,8 @@ from .status import StatusTrait from .volume import SoundVolumeTrait +_LOGGER = logging.getLogger(__name__) + __all__ = [ "create", "PropertiesApi", @@ -40,10 +43,12 @@ class PropertiesApi(Trait): # In the future optional fields can be added below based on supported features - def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel) -> None: + def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel) -> None: """Initialize the V1TraitProps with None values.""" self.status = StatusTrait(product) + #self.status._rpc_channel = rpc_channel self.maps = MapsTrait(self.status) + #self.maps._rpc_channel = mqtt_rpc_channel # This is a hack to allow setting the rpc_channel on all traits. This is # used so we can preserve the dataclass behavior when the values in the @@ -53,9 +58,12 @@ def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel) -> None: if (trait := getattr(self, item.name, None)) is None: trait = item.type() setattr(self, item.name, trait) - trait._rpc_channel = rpc_channel + if hasattr(trait, "mqtt_rpc_channel"): # @common.mqtt_rpc_channel + trait._rpc_channel = mqtt_rpc_channel + else: + trait._rpc_channel = rpc_channel -def create(product: HomeDataProduct, rpc_channel: V1RpcChannel) -> PropertiesApi: +def create(product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel) -> PropertiesApi: """Create traits for V1 devices.""" - return PropertiesApi(product, rpc_channel) + return PropertiesApi(product, rpc_channel, mqtt_rpc_channel) diff --git a/roborock/devices/traits/v1/common.py b/roborock/devices/traits/v1/common.py index 7b14d0d3..df6144c3 100644 --- a/roborock/devices/traits/v1/common.py +++ b/roborock/devices/traits/v1/common.py @@ -5,12 +5,16 @@ from abc import ABC from dataclasses import dataclass, fields +from functools import wraps +import logging from typing import ClassVar, Self from roborock.containers import RoborockBase from roborock.devices.v1_rpc_channel import V1RpcChannel from roborock.roborock_typing import RoborockCommand +_LOGGER = logging.getLogger(__name__) + V1ResponseData = dict | list | int | str @@ -115,3 +119,16 @@ def _parse_response(cls, response: V1ResponseData) -> Self: raise ValueError(f"Unexpected response format: {response!r}") value_field = _get_value_field(cls) return cls(**{value_field: response}) + + +def mqtt_rpc_channel(cls): + """Decorator to mark a function as cloud only. + + Normally a trait uses an adaptive rpc channel that can use either local + or cloud communication depending on what is available. This will force + the trait to always use the cloud rpc channel. + """ + def wrapper(*args, **kwargs): + return cls(*args, **kwargs) + cls.mqtt_rpc_channel = True # type: ignore[attr-defined] + return wrapper diff --git a/roborock/devices/traits/v1/maps.py b/roborock/devices/traits/v1/maps.py index 2da32de0..bfd78233 100644 --- a/roborock/devices/traits/v1/maps.py +++ b/roborock/devices/traits/v1/maps.py @@ -16,6 +16,7 @@ _LOGGER = logging.getLogger(__name__) +@common.mqtt_rpc_channel class MapsTrait(MultiMapsList, common.V1TraitMixin): """Trait for managing the maps of Roborock devices. diff --git a/tests/devices/test_v1_device.py b/tests/devices/test_v1_device.py index e164f8bc..66207673 100644 --- a/tests/devices/test_v1_device.py +++ b/tests/devices/test_v1_device.py @@ -35,13 +35,18 @@ def rpc_channel_fixture() -> AsyncMock: return AsyncMock() +@pytest.fixture(autouse=True, name="mqtt_rpc_channel") +def mqtt_rpc_channel_fixture() -> AsyncMock: + """Fixture to set up the channel for tests.""" + return AsyncMock() + @pytest.fixture(autouse=True, name="device") -def device_fixture(channel: AsyncMock, rpc_channel: AsyncMock) -> RoborockDevice: +def device_fixture(channel: AsyncMock, rpc_channel: AsyncMock, mqtt_rpc_channel: AsyncMock) -> RoborockDevice: """Fixture to set up the device for tests.""" return RoborockDevice( device_info=HOME_DATA.devices[0], channel=channel, - trait=v1.create(HOME_DATA.products[0], rpc_channel), + trait=v1.create(HOME_DATA.products[0], rpc_channel, mqtt_rpc_channel), ) diff --git a/tests/devices/traits/v1/fixtures.py b/tests/devices/traits/v1/fixtures.py index 188d8a9a..58f5f218 100644 --- a/tests/devices/traits/v1/fixtures.py +++ b/tests/devices/traits/v1/fixtures.py @@ -27,11 +27,17 @@ def rpc_channel_fixture() -> AsyncMock: return AsyncMock() + +@pytest.fixture(autouse=True, name="mock_mqtt_rpc_channel") +def mqtt_rpc_channel_fixture() -> AsyncMock: + """Fixture to set up the channel for tests.""" + return AsyncMock() + @pytest.fixture(autouse=True, name="device") -def device_fixture(channel: AsyncMock, mock_rpc_channel: AsyncMock) -> RoborockDevice: +def device_fixture(channel: AsyncMock, mock_rpc_channel: AsyncMock, mock_mqtt_rpc_channel: AsyncMock) -> RoborockDevice: """Fixture to set up the device for tests.""" return RoborockDevice( device_info=HOME_DATA.devices[0], channel=channel, - trait=v1.create(HOME_DATA.products[0], mock_rpc_channel), + trait=v1.create(HOME_DATA.products[0], mock_rpc_channel, mock_mqtt_rpc_channel), ) diff --git a/tests/devices/traits/v1/test_maps.py b/tests/devices/traits/v1/test_maps.py index 02caef0d..660acc3f 100644 --- a/tests/devices/traits/v1/test_maps.py +++ b/tests/devices/traits/v1/test_maps.py @@ -58,15 +58,20 @@ def maps_trait(device: RoborockDevice) -> MapsTrait: async def test_refresh_maps_trait( maps_trait: MapsTrait, mock_rpc_channel: AsyncMock, + mock_mqtt_rpc_channel: AsyncMock, status_trait: StatusTrait, ) -> None: """Test successfully getting multi maps list.""" # Setup mock to return the sample multi maps list mock_rpc_channel.send_command.side_effect = [ mock_data.STATUS, # Initial status fetch + ] + mock_mqtt_rpc_channel.send_command.side_effect = [ MULTI_MAP_LIST_DATA, ] await status_trait.refresh() + assert status_trait.current_map == 0 + # Populating the status information gives us the current map # flag, but we have not loaded the rest of the information. assert maps_trait.current_map == 0 @@ -96,22 +101,27 @@ async def test_refresh_maps_trait( assert maps_trait.current_map_info.name == "Map 1" # Verify the RPC call was made correctly - assert mock_rpc_channel.send_command.call_count == 2 + assert mock_rpc_channel.send_command.call_count == 1 mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_STATUS) - mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_MULTI_MAPS_LIST) + assert mock_mqtt_rpc_channel.send_command.call_count == 1 + mock_mqtt_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_MULTI_MAPS_LIST) async def test_set_current_map( status_trait: StatusTrait, maps_trait: MapsTrait, mock_rpc_channel: AsyncMock, + mock_mqtt_rpc_channel: AsyncMock, ) -> None: """Test successfully setting the current map.""" + assert hasattr(maps_trait, "mqtt_rpc_channel") mock_rpc_channel.send_command.side_effect = [ mock_data.STATUS, # Initial status fetch + UPDATED_STATUS, # Response for refreshing status + ] + mock_mqtt_rpc_channel.send_command.side_effect = [ MULTI_MAP_LIST_DATA, # Response for LOAD_MULTI_MAP {}, # Response for setting the current map - UPDATED_STATUS, # Response for refreshing status ] await status_trait.refresh() @@ -139,7 +149,8 @@ async def test_set_current_map( # 2. GET_MULTI_MAPS_LIST to get the map list # 3. LOAD_MULTI_MAP to set the map # 4. GET_STATUS to refresh the current map in status - assert mock_rpc_channel.send_command.call_count == 4 + assert mock_rpc_channel.send_command.call_count == 2 mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_STATUS) - mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_MULTI_MAPS_LIST) - mock_rpc_channel.send_command.assert_any_call(RoborockCommand.LOAD_MULTI_MAP, params=[123]) + assert mock_mqtt_rpc_channel.send_command.call_count == 2 + mock_mqtt_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_MULTI_MAPS_LIST) + mock_mqtt_rpc_channel.send_command.assert_any_call(RoborockCommand.LOAD_MULTI_MAP, params=[123]) From 002317477e2fc267ce67e17cb494843f28997928 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sat, 27 Sep 2025 22:12:57 -0700 Subject: [PATCH 5/8] chore: fix lint errors --- roborock/devices/traits/v1/__init__.py | 6 +++--- roborock/devices/traits/v1/common.py | 7 ++++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/roborock/devices/traits/v1/__init__.py b/roborock/devices/traits/v1/__init__.py index f1811af5..c9b1ccc0 100644 --- a/roborock/devices/traits/v1/__init__.py +++ b/roborock/devices/traits/v1/__init__.py @@ -1,7 +1,7 @@ """Create traits for V1 devices.""" -from dataclasses import dataclass, field, fields import logging +from dataclasses import dataclass, field, fields from roborock.containers import HomeData, HomeDataProduct from roborock.devices.traits import Trait @@ -46,9 +46,9 @@ class PropertiesApi(Trait): def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel) -> None: """Initialize the V1TraitProps with None values.""" self.status = StatusTrait(product) - #self.status._rpc_channel = rpc_channel + # self.status._rpc_channel = rpc_channel self.maps = MapsTrait(self.status) - #self.maps._rpc_channel = mqtt_rpc_channel + # self.maps._rpc_channel = mqtt_rpc_channel # This is a hack to allow setting the rpc_channel on all traits. This is # used so we can preserve the dataclass behavior when the values in the diff --git a/roborock/devices/traits/v1/common.py b/roborock/devices/traits/v1/common.py index df6144c3..fc13838f 100644 --- a/roborock/devices/traits/v1/common.py +++ b/roborock/devices/traits/v1/common.py @@ -3,10 +3,9 @@ This is an internal library and should not be used directly by consumers. """ +import logging from abc import ABC from dataclasses import dataclass, fields -from functools import wraps -import logging from typing import ClassVar, Self from roborock.containers import RoborockBase @@ -123,12 +122,14 @@ def _parse_response(cls, response: V1ResponseData) -> Self: def mqtt_rpc_channel(cls): """Decorator to mark a function as cloud only. - + Normally a trait uses an adaptive rpc channel that can use either local or cloud communication depending on what is available. This will force the trait to always use the cloud rpc channel. """ + def wrapper(*args, **kwargs): return cls(*args, **kwargs) + cls.mqtt_rpc_channel = True # type: ignore[attr-defined] return wrapper From 3810cb60446466ec66e636b123888013eb4a4f46 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sat, 27 Sep 2025 22:13:24 -0700 Subject: [PATCH 6/8] chore: fix lint errors --- tests/devices/test_v1_device.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/devices/test_v1_device.py b/tests/devices/test_v1_device.py index 66207673..30515b48 100644 --- a/tests/devices/test_v1_device.py +++ b/tests/devices/test_v1_device.py @@ -40,6 +40,7 @@ def mqtt_rpc_channel_fixture() -> AsyncMock: """Fixture to set up the channel for tests.""" return AsyncMock() + @pytest.fixture(autouse=True, name="device") def device_fixture(channel: AsyncMock, rpc_channel: AsyncMock, mqtt_rpc_channel: AsyncMock) -> RoborockDevice: """Fixture to set up the device for tests.""" From 01026ab29015c98fa6c11b5fbd545c958b3b4d03 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sat, 27 Sep 2025 23:18:34 -0700 Subject: [PATCH 7/8] chore: fix lint errors --- tests/devices/traits/v1/fixtures.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/devices/traits/v1/fixtures.py b/tests/devices/traits/v1/fixtures.py index 58f5f218..8def9231 100644 --- a/tests/devices/traits/v1/fixtures.py +++ b/tests/devices/traits/v1/fixtures.py @@ -27,12 +27,12 @@ def rpc_channel_fixture() -> AsyncMock: return AsyncMock() - @pytest.fixture(autouse=True, name="mock_mqtt_rpc_channel") def mqtt_rpc_channel_fixture() -> AsyncMock: """Fixture to set up the channel for tests.""" return AsyncMock() + @pytest.fixture(autouse=True, name="device") def device_fixture(channel: AsyncMock, mock_rpc_channel: AsyncMock, mock_mqtt_rpc_channel: AsyncMock) -> RoborockDevice: """Fixture to set up the device for tests.""" From d56c5fe72176e22116d422ded49b23048e2a15fc Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sun, 28 Sep 2025 18:30:53 -0700 Subject: [PATCH 8/8] chore: add comment describing the decorator check --- roborock/devices/traits/v1/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/roborock/devices/traits/v1/__init__.py b/roborock/devices/traits/v1/__init__.py index c9b1ccc0..d152b01a 100644 --- a/roborock/devices/traits/v1/__init__.py +++ b/roborock/devices/traits/v1/__init__.py @@ -46,9 +46,7 @@ class PropertiesApi(Trait): def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel) -> None: """Initialize the V1TraitProps with None values.""" self.status = StatusTrait(product) - # self.status._rpc_channel = rpc_channel self.maps = MapsTrait(self.status) - # self.maps._rpc_channel = mqtt_rpc_channel # This is a hack to allow setting the rpc_channel on all traits. This is # used so we can preserve the dataclass behavior when the values in the @@ -58,7 +56,9 @@ def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc if (trait := getattr(self, item.name, None)) is None: trait = item.type() setattr(self, item.name, trait) - if hasattr(trait, "mqtt_rpc_channel"): # @common.mqtt_rpc_channel + # The decorator `@common.mqtt_rpc_channel` means that the trait needs + # to use the mqtt_rpc_channel (cloud only) instead of the rpc_channel (adaptive) + if hasattr(trait, "mqtt_rpc_channel"): trait._rpc_channel = mqtt_rpc_channel else: trait._rpc_channel = rpc_channel