feat: Add v1 rooms support to the device traits API

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():

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)

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)

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

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

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),
     )
 
 

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),
     )

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)

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
 # ---

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\"]]}"}}

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]]}"}}