diff --git a/roborock/cli.py b/roborock/cli.py index d98582df..8ceb1d51 100644 --- a/roborock/cli.py +++ b/roborock/cli.py @@ -475,6 +475,16 @@ async def reset_consumable(ctx, device_id: str, consumable: str): click.echo(f"Reset {consumable} for device {device_id}") +@session.command() +@click.option("--device_id", required=True) +@click.pass_context +@async_command +async def rooms(ctx, device_id: str): + """Get device room mapping info.""" + context: RoborockContext = ctx.obj + await _display_v1_trait(context, device_id, lambda v1: v1.rooms) + + @click.command() @click.option("--device_id", required=True) @click.option("--cmd", required=True) @@ -719,6 +729,7 @@ def write_markdown_table(product_features: dict[str, dict[str, any]], all_featur cli.add_command(maps) cli.add_command(consumables) cli.add_command(reset_consumable) +cli.add_command(rooms) def main(): diff --git a/roborock/devices/device_manager.py b/roborock/devices/device_manager.py index 0361c8cc..b01bb671 100644 --- a/roborock/devices/device_manager.py +++ b/roborock/devices/device_manager.py @@ -35,7 +35,7 @@ HomeDataApi = Callable[[], Awaitable[HomeData]] -DeviceCreator = Callable[[HomeDataDevice, HomeDataProduct], RoborockDevice] +DeviceCreator = Callable[[HomeData, HomeDataDevice, HomeDataProduct], RoborockDevice] class DeviceVersion(enum.StrEnum): @@ -84,7 +84,7 @@ async def discover_devices(self) -> list[RoborockDevice]: for duid, (device, product) in device_products.items(): if duid in self._devices: continue - new_device = self._device_creator(device, product) + new_device = self._device_creator(home_data, device, product) await new_device.connect() new_devices[duid] = new_device @@ -143,13 +143,13 @@ async def create_device_manager( mqtt_params = create_mqtt_params(user_data.rriot) mqtt_session = await create_lazy_mqtt_session(mqtt_params) - def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice: + def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice: channel: Channel trait: Trait 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, channel.mqtt_rpc_channel) + trait = v1.create(product, home_data, 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 7edf5f43..065d4744 100644 --- a/roborock/devices/traits/v1/__init__.py +++ b/roborock/devices/traits/v1/__init__.py @@ -12,6 +12,7 @@ from .consumeable import ConsumableTrait from .do_not_disturb import DoNotDisturbTrait from .maps import MapsTrait +from .rooms import RoomsTrait from .status import StatusTrait from .volume import SoundVolumeTrait @@ -41,14 +42,18 @@ class PropertiesApi(Trait): dnd: DoNotDisturbTrait clean_summary: CleanSummaryTrait sound_volume: SoundVolumeTrait + rooms: RoomsTrait maps: MapsTrait consumables: ConsumableTrait # In the future optional fields can be added below based on supported features - def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel) -> None: - """Initialize the V1TraitProps with None values.""" + def __init__( + self, product: HomeDataProduct, home_data: HomeData, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel + ) -> None: + """Initialize the V1TraitProps.""" self.status = StatusTrait(product) + self.rooms = RoomsTrait(home_data) self.maps = MapsTrait(self.status) # This is a hack to allow setting the rpc_channel on all traits. This is @@ -67,6 +72,8 @@ def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc trait._rpc_channel = rpc_channel -def create(product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel) -> PropertiesApi: +def create( + product: HomeDataProduct, home_data: HomeData, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel +) -> PropertiesApi: """Create traits for V1 devices.""" - return PropertiesApi(product, rpc_channel, mqtt_rpc_channel) + return PropertiesApi(product, home_data, rpc_channel, mqtt_rpc_channel) diff --git a/roborock/devices/traits/v1/common.py b/roborock/devices/traits/v1/common.py index fc13838f..e02ec873 100644 --- a/roborock/devices/traits/v1/common.py +++ b/roborock/devices/traits/v1/common.py @@ -38,7 +38,7 @@ class V1TraitMixin(ABC): command: ClassVar[RoborockCommand] @classmethod - def _parse_type_response(cls, response: V1ResponseData) -> Self: + def _parse_type_response(cls, response: V1ResponseData) -> RoborockBase: """Parse the response from the device into a a RoborockBase. Subclasses should override this method to implement custom parsing @@ -53,7 +53,7 @@ def _parse_type_response(cls, response: V1ResponseData) -> Self: raise ValueError(f"Unexpected {cls} response format: {response!r}") return cls.from_dict(response) - def _parse_response(self, response: V1ResponseData) -> Self: + def _parse_response(self, response: V1ResponseData) -> RoborockBase: """Parse the response from the device into a a RoborockBase. This is used by subclasses that want to override the class diff --git a/roborock/devices/traits/v1/rooms.py b/roborock/devices/traits/v1/rooms.py new file mode 100644 index 00000000..b4b36f92 --- /dev/null +++ b/roborock/devices/traits/v1/rooms.py @@ -0,0 +1,93 @@ +"""Trait for managing room mappings on Roborock devices.""" + +import logging +from dataclasses import dataclass + +from roborock.containers import HomeData, RoborockBase, RoomMapping +from roborock.devices.traits.v1 import common +from roborock.roborock_typing import RoborockCommand + +_LOGGER = logging.getLogger(__name__) + +_DEFAULT_NAME = "Unknown" + + +@dataclass +class NamedRoomMapping(RoomMapping): + """Dataclass representing a mapping of a room segment to a name. + + The name information is not provided by the device directly, but is provided + from the HomeData based on the iot_id from the room. + """ + + name: str + """The human-readable name of the room, if available.""" + + +@dataclass +class Rooms(RoborockBase): + """Dataclass representing a collection of room mappings.""" + + rooms: list[NamedRoomMapping] | None = None + """List of room mappings.""" + + @property + def room_map(self) -> dict[int, NamedRoomMapping]: + """Returns a mapping of segment_id to NamedRoomMapping.""" + if self.rooms is None: + return {} + return {room.segment_id: room for room in self.rooms} + + +class RoomsTrait(Rooms, common.V1TraitMixin): + """Trait for managing the room mappings of Roborock devices.""" + + command = RoborockCommand.GET_ROOM_MAPPING + + def __init__(self, home_data: HomeData) -> None: + """Initialize the RoomsTrait.""" + super().__init__() + self._home_data = home_data + + @property + def _iot_id_room_name_map(self) -> dict[str, str]: + """Returns a dictionary of Room IOT IDs to room names.""" + return {str(room.id): room.name for room in self._home_data.rooms or ()} + + def _parse_response(self, response: common.V1ResponseData) -> Rooms: + """Parse the response from the device into a list of NamedRoomMapping.""" + if not isinstance(response, list): + raise ValueError(f"Unexpected RoomsTrait response format: {response!r}") + name_map = self._iot_id_room_name_map + segment_pairs = _extract_segment_pairs(response) + return Rooms( + rooms=[ + NamedRoomMapping(segment_id=segment_id, iot_id=iot_id, name=name_map.get(iot_id, _DEFAULT_NAME)) + for segment_id, iot_id in segment_pairs + ] + ) + + +def _extract_segment_pairs(response: list) -> list[tuple[int, str]]: + """Extract segment_id and iot_id pairs from the response. + + The response format can be either a flat list of [segment_id, iot_id] or a + list of lists, where each inner list is a pair of [segment_id, iot_id]. This + function normalizes the response into a list of (segment_id, iot_id) tuples + + NOTE: We currently only partial samples of the room mapping formats, so + improving test coverage with samples from a real device with this format + would be helpful. + """ + if len(response) == 2 and not isinstance(response[0], list): + segment_id, iot_id = response[0], response[1] + return [(segment_id, iot_id)] + + segment_pairs: list[tuple[int, str]] = [] + for part in response: + if not isinstance(part, list) or len(part) < 2: + _LOGGER.warning("Unexpected room mapping entry format: %r", part) + continue + segment_id, iot_id = part[0], part[1] + segment_pairs.append((segment_id, iot_id)) + return segment_pairs diff --git a/tests/devices/test_v1_device.py b/tests/devices/test_v1_device.py index 30515b48..3628772f 100644 --- a/tests/devices/test_v1_device.py +++ b/tests/devices/test_v1_device.py @@ -47,7 +47,7 @@ def device_fixture(channel: AsyncMock, rpc_channel: AsyncMock, mqtt_rpc_channel: return RoborockDevice( device_info=HOME_DATA.devices[0], channel=channel, - trait=v1.create(HOME_DATA.products[0], rpc_channel, mqtt_rpc_channel), + trait=v1.create(HOME_DATA.products[0], HOME_DATA, rpc_channel, mqtt_rpc_channel), ) diff --git a/tests/devices/traits/v1/fixtures.py b/tests/devices/traits/v1/fixtures.py index 8def9231..431741e0 100644 --- a/tests/devices/traits/v1/fixtures.py +++ b/tests/devices/traits/v1/fixtures.py @@ -39,5 +39,5 @@ def device_fixture(channel: AsyncMock, mock_rpc_channel: AsyncMock, mock_mqtt_rp return RoborockDevice( device_info=HOME_DATA.devices[0], channel=channel, - trait=v1.create(HOME_DATA.products[0], mock_rpc_channel, mock_mqtt_rpc_channel), + trait=v1.create(HOME_DATA.products[0], HOME_DATA, mock_rpc_channel, mock_mqtt_rpc_channel), ) diff --git a/tests/devices/traits/v1/test_rooms.py b/tests/devices/traits/v1/test_rooms.py new file mode 100644 index 00000000..271ecd61 --- /dev/null +++ b/tests/devices/traits/v1/test_rooms.py @@ -0,0 +1,72 @@ +"""Tests for the RoomMapping related functionality.""" + +from typing import Any +from unittest.mock import AsyncMock + +import pytest + +from roborock.devices.device import RoborockDevice +from roborock.devices.traits.v1.rooms import RoomsTrait +from roborock.devices.traits.v1.status import StatusTrait +from roborock.roborock_typing import RoborockCommand + + +@pytest.fixture +def status_trait(device: RoborockDevice) -> StatusTrait: + """Create a StatusTrait instance with mocked dependencies.""" + assert device.v1_properties + return device.v1_properties.status + + +@pytest.fixture +def rooms_trait(device: RoborockDevice) -> RoomsTrait: + """Create a RoomsTrait instance with mocked dependencies.""" + assert device.v1_properties + return device.v1_properties.rooms + + +# Rooms from mock_data.HOME_DATA +# {"id": 2362048, "name": "Example room 1"}, +# {"id": 2362044, "name": "Example room 2"}, +# {"id": 2362041, "name": "Example room 3"}, +@pytest.mark.parametrize( + ("room_mapping_data"), + [ + ([[16, "2362048"], [17, "2362044"], [18, "2362041"]]), + ([[16, "2362048", 6], [17, "2362044", 14], [18, "2362041", 13]]), + ], +) +async def test_refresh_rooms_trait( + rooms_trait: RoomsTrait, + mock_rpc_channel: AsyncMock, + room_mapping_data: list[Any], +) -> None: + """Test successfully getting room mapping.""" + # Setup mock to return the sample room mapping + mock_rpc_channel.send_command.side_effect = [room_mapping_data] + # Before refresh, rooms should be empty + assert not rooms_trait.rooms + + # Load the room mapping information + refreshed_trait = await rooms_trait.refresh() + + # Verify the room mappings are now populated + assert refreshed_trait.rooms + rooms = refreshed_trait.rooms + assert len(rooms) == 3 + + assert rooms[0].segment_id == 16 + assert rooms[0].name == "Example room 1" + assert rooms[0].iot_id == "2362048" + + assert rooms[1].segment_id == 17 + assert rooms[1].name == "Example room 2" + assert rooms[1].iot_id == "2362044" + + assert rooms[2].segment_id == 18 + assert rooms[2].name == "Example room 3" + assert rooms[2].iot_id == "2362041" + + # Verify the RPC call was made correctly + assert mock_rpc_channel.send_command.call_count == 1 + mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_ROOM_MAPPING) diff --git a/tests/protocols/__snapshots__/test_v1_protocol.ambr b/tests/protocols/__snapshots__/test_v1_protocol.ambr index 0b736df6..95337182 100644 --- a/tests/protocols/__snapshots__/test_v1_protocol.ambr +++ b/tests/protocols/__snapshots__/test_v1_protocol.ambr @@ -93,6 +93,71 @@ ] ''' # --- +# name: test_decode_rpc_payload[get_room_mapping2] + 20001 +# --- +# name: test_decode_rpc_payload[get_room_mapping2].1 + ''' + [ + [ + 16, + "2537178", + 6 + ], + [ + 17, + "2537175", + 14 + ], + [ + 18, + "2537174", + 13 + ], + [ + 19, + "2537176", + 14 + ], + [ + 20, + "10655627", + 12 + ], + [ + 21, + "2537145", + 2 + ], + [ + 22, + "2537147", + 12 + ] + ] + ''' +# --- +# name: test_decode_rpc_payload[get_room_mapping] + 20001 +# --- +# name: test_decode_rpc_payload[get_room_mapping].1 + ''' + [ + [ + 16, + "3031886" + ], + [ + 17, + "3031880" + ], + [ + 18, + "3031883" + ] + ] + ''' +# --- # name: test_decode_rpc_payload[get_status] 20001 # --- diff --git a/tests/protocols/testdata/v1_protocol/get_room_mapping.json b/tests/protocols/testdata/v1_protocol/get_room_mapping.json new file mode 100644 index 00000000..ef363f0c --- /dev/null +++ b/tests/protocols/testdata/v1_protocol/get_room_mapping.json @@ -0,0 +1 @@ +{"t":1759590351,"dps":{"102":"{\"id\":20001,\"result\":[[16,\"3031886\"],[17,\"3031880\"],[18,\"3031883\"]]}"}} diff --git a/tests/protocols/testdata/v1_protocol/get_room_mapping2.json b/tests/protocols/testdata/v1_protocol/get_room_mapping2.json new file mode 100644 index 00000000..49fd1195 --- /dev/null +++ b/tests/protocols/testdata/v1_protocol/get_room_mapping2.json @@ -0,0 +1 @@ +{"t":1759590351,"dps":{"102":"{\"id\":20001,\"result\":[[16, \"2537178\", 6], [17, \"2537175\", 14], [18, \"2537174\", 13], [19, \"2537176\", 14], [20, \"10655627\", 12], [21, \"2537145\", 2], [22, \"2537147\", 12]]}"}}