address allens comments

Author: Lash-L

roborock/devices/traits/v1/rooms.py

@@ -2,7 +2,6 @@
 
 import logging
 from dataclasses import dataclass
-from functools import cached_property
 
 from roborock.data import HomeData, HomeDataRoom, NamedRoomMapping, RoborockBase
 from roborock.devices.traits.v1 import common
@@ -90,17 +89,16 @@ def __init__(self, home_data: HomeData, device_uid: str, web_api: UserWebApiClie
         super().__init__()
         self._home_data = home_data
         self._device_uid = device_uid
+        self._shared_device_uid = next(
+            (device.duid for device in home_data.received_devices if device.duid == device_uid), None
+        )
         self._web_api = web_api
         self._discovered_iot_ids: set[str] = set()
-        self._shared_room_names: dict[str, str] = {}
-
-    @cached_property
-    def _is_shared(self) -> bool:
-        return any(d.duid == self._device_uid for d in self._home_data.received_devices)
+        self._room_names: dict[str, str] = dict(home_data.rooms_name_map)
 
     @property
     def _room_name_map(self) -> dict[str, str]:
-        return {**self._home_data.rooms_name_map, **self._shared_room_names}
+        return self._room_names
 
     async def refresh(self) -> None:
         """Refresh room mappings and backfill unknown room names from the web API."""
@@ -120,9 +118,8 @@ async def refresh(self) -> None:
             _LOGGER.debug("Refreshing room list to discover new room names")
             if updated_rooms := await self._refresh_rooms():
                 _LOGGER.debug("Updating rooms: %s", list(updated_rooms))
-                if self._is_shared:
-                    self._shared_room_names = {room.iot_id: room.name for room in updated_rooms}
-                else:
+                self._room_names = {room.iot_id: room.name for room in updated_rooms}
+                if self._shared_device_uid is None:
                     self._home_data.rooms = updated_rooms
             self._discovered_iot_ids.update(new_iot_ids)
         try:
@@ -141,8 +138,11 @@ async def refresh(self) -> None:
     async def _refresh_rooms(self) -> list[HomeDataRoom]:
         """Fetch the latest rooms from the web API."""
         try:
-            if self._is_shared:
-                return await self._web_api.get_shared_device_rooms(self._device_uid)
+            if self._shared_device_uid is not None:
+                rooms_by_id = {room.iot_id: room for room in self._home_data.rooms}
+                shared_rooms = await self._web_api.get_shared_device_rooms(self._shared_device_uid)
+                rooms_by_id.update({room.iot_id: room for room in shared_rooms})
+                return list(rooms_by_id.values())
             return await self._web_api.get_rooms()
         except Exception:
             _LOGGER.debug("Failed to fetch rooms from web API", exc_info=True)

tests/devices/traits/v1/fixtures.py

@@ -46,22 +46,28 @@ def web_api_client_fixture() -> AsyncMock:
     return AsyncMock()
 
 
+@pytest.fixture(name="home_data")
+def home_data_fixture(request: pytest.FixtureRequest) -> HomeData:
+    """Fixture to provide HomeData, optionally overridden via indirect parametrization."""
+    return deepcopy(getattr(request, "param", HOME_DATA))
+
+
 @pytest.fixture(autouse=True, name="roborock_cache")
 def roborock_cache_fixture() -> Cache:
     """Fixture to provide a NoCache instance for tests."""
     return InMemoryCache()
 
 
 @pytest.fixture(autouse=True, name="device_cache")
-def device_cache_fixture(roborock_cache: Cache) -> DeviceCache:
+def device_cache_fixture(roborock_cache: Cache, home_data: HomeData) -> DeviceCache:
     """Fixture to provide a DeviceCache instance for tests."""
-    return DeviceCache(HOME_DATA.devices[0].duid, roborock_cache)
+    return DeviceCache(home_data.get_all_devices()[0].duid, roborock_cache)
 
 
 @pytest.fixture(name="device_info")
-def device_info_fixture() -> HomeDataDevice:
+def device_info_fixture(home_data: HomeData) -> HomeDataDevice:
     """Fixture to provide a DeviceInfo instance for tests."""
-    return HOME_DATA.devices[0]
+    return home_data.get_all_devices()[0]
 
 
 @pytest.fixture(name="products")
@@ -79,6 +85,7 @@ def device_fixture(
     web_api_client: AsyncMock,
     device_cache: DeviceCache,
     device_info: HomeDataDevice,
+    home_data: HomeData,
     products: list[HomeDataProduct],
 ) -> RoborockDevice:
     """Fixture to set up the device for tests."""
@@ -90,7 +97,7 @@ def device_fixture(
         trait=v1.create(
             device_info.duid,
             product,
-            deepcopy(HOME_DATA),
+            home_data,
             mock_rpc_channel,
             mock_mqtt_rpc_channel,
             mock_map_rpc_channel,

tests/devices/traits/v1/test_rooms.py

@@ -1,15 +1,17 @@
 """Tests for the RoomMapping related functionality."""
 
+from copy import deepcopy
 from typing import Any
 from unittest.mock import AsyncMock
 
 import pytest
 
-from roborock.data.containers import HomeDataRoom, NamedRoomMapping
+from roborock.data.containers import HomeData, HomeDataRoom, NamedRoomMapping
 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
+from tests.devices.traits.v1.fixtures import HOME_DATA
 
 
 @pytest.fixture
@@ -79,29 +81,25 @@ async def test_refresh_unknown_room_names_overwrites_home_data(
     mock_rpc_channel: AsyncMock,
 ) -> None:
     """Test web rooms are used to refresh home data for missing iot ids."""
-    original_rooms = list(rooms_trait._home_data.rooms or ())
-    try:
-        web_api_client.get_rooms.return_value = [
-            HomeDataRoom(id=2362048, name="Living Room"),
-            HomeDataRoom(id=2362044, name="Example room 2"),
-            HomeDataRoom(id=2362041, name="Example room 3"),
-            HomeDataRoom(id=9999999, name="Office"),
-        ]
-
-        room_mapping_data = [[16, "2362048"], [17, "9999999"]]
-        mock_rpc_channel.send_command.side_effect = [room_mapping_data]
+    web_api_client.get_rooms.return_value = [
+        HomeDataRoom(id=2362048, name="Living Room"),
+        HomeDataRoom(id=2362044, name="Example room 2"),
+        HomeDataRoom(id=2362041, name="Example room 3"),
+        HomeDataRoom(id=9999999, name="Office"),
+    ]
+
+    room_mapping_data = [[16, "2362048"], [17, "9999999"]]
+    mock_rpc_channel.send_command.side_effect = [room_mapping_data]
 
-        await rooms_trait.refresh()
+    await rooms_trait.refresh()
 
-        assert rooms_trait.rooms
-        assert rooms_trait.rooms[0] == NamedRoomMapping(segment_id=16, iot_id="2362048", raw_name="Living Room")
-        assert rooms_trait.rooms[1] == NamedRoomMapping(segment_id=17, iot_id="9999999", raw_name="Office")
+    assert rooms_trait.rooms
+    assert rooms_trait.rooms[0] == NamedRoomMapping(segment_id=16, iot_id="2362048", raw_name="Living Room")
+    assert rooms_trait.rooms[1] == NamedRoomMapping(segment_id=17, iot_id="9999999", raw_name="Office")
 
-        home_data_rooms = {str(room.id): room.name for room in rooms_trait._home_data.rooms or ()}
-        assert home_data_rooms["2362048"] == "Living Room"
-        assert home_data_rooms["9999999"] == "Office"
-    finally:
-        rooms_trait._home_data.rooms = original_rooms
+    home_data_rooms = {str(room.id): room.name for room in rooms_trait._home_data.rooms or ()}
+    assert home_data_rooms["2362048"] == "Living Room"
+    assert home_data_rooms["9999999"] == "Office"
 
 
 async def test_refresh_unknown_room_names_web_api_called_once(
@@ -110,25 +108,21 @@ async def test_refresh_unknown_room_names_web_api_called_once(
     mock_rpc_channel: AsyncMock,
 ) -> None:
     """Test unknown room IDs trigger one web lookup per iot_id."""
-    original_rooms = list(rooms_trait._home_data.rooms or ())
-    try:
-        web_api_client.get_rooms.return_value = [
-            HomeDataRoom(id=9999911, name="Living Room"),
-        ]
+    web_api_client.get_rooms.return_value = [
+        HomeDataRoom(id=9999911, name="Living Room"),
+    ]
 
-        room_mapping_data = [[16, "9999911"]]
-        mock_rpc_channel.send_command.side_effect = [room_mapping_data, room_mapping_data]
+    room_mapping_data = [[16, "9999911"]]
+    mock_rpc_channel.send_command.side_effect = [room_mapping_data, room_mapping_data]
 
-        await rooms_trait.refresh()
-        assert rooms_trait.rooms
-        assert rooms_trait.rooms[0].name == "Living Room"
+    await rooms_trait.refresh()
+    assert rooms_trait.rooms
+    assert rooms_trait.rooms[0].name == "Living Room"
 
-        await rooms_trait.refresh()
-        assert rooms_trait.rooms
-        assert rooms_trait.rooms[0].name == "Living Room"
-        web_api_client.get_rooms.assert_called_once()
-    finally:
-        rooms_trait._home_data.rooms = original_rooms
+    await rooms_trait.refresh()
+    assert rooms_trait.rooms
+    assert rooms_trait.rooms[0].name == "Living Room"
+    web_api_client.get_rooms.assert_called_once()
 
 
 async def test_refresh_unknown_room_names_unresolved_uses_room_fallback(
@@ -209,32 +203,34 @@ async def test_refresh_unknown_room_names_failure_falls_back_to_room_segment_id(
     web_api_client.get_rooms.assert_called_once()
 
 
-async def test_refresh_shared_room_names_use_shared_device_rooms_without_mutating_home_data(
+def _build_shared_home_data() -> HomeData:
+    home_data = deepcopy(HOME_DATA)
+    home_data.received_devices = [home_data.devices.pop(0)]
+    return home_data
+
+
+@pytest.mark.parametrize("home_data", [_build_shared_home_data()], indirect=True)
+async def test_refresh_shared_room_names_use_shared_device_rooms(
     rooms_trait: RoomsTrait,
+    home_data: HomeData,
     web_api_client: AsyncMock,
     mock_rpc_channel: AsyncMock,
 ) -> None:
     """Test shared devices resolve room names via the shared-device room list."""
-    original_rooms = list(rooms_trait._home_data.rooms or ())
-    try:
-        # Mark the device as shared by adding it to received_devices
-        device = next(d for d in rooms_trait._home_data.devices if d.duid == rooms_trait._device_uid)
-        rooms_trait._home_data.received_devices = [device]
-        rooms_trait._home_data.devices = []
-
-        web_api_client.get_shared_device_rooms.return_value = [
-            HomeDataRoom(id=9999999, name="Office"),
-        ]
-        room_mapping_data = [[16, "2362048"], [17, "9999999"]]
-        mock_rpc_channel.send_command.side_effect = [room_mapping_data]
-
-        await rooms_trait.refresh()
-
-        assert rooms_trait.rooms
-        assert rooms_trait.rooms[0] == NamedRoomMapping(segment_id=16, iot_id="2362048", raw_name="Example room 1")
-        assert rooms_trait.rooms[1] == NamedRoomMapping(segment_id=17, iot_id="9999999", raw_name="Office")
-        assert rooms_trait._home_data.rooms == original_rooms
-        web_api_client.get_shared_device_rooms.assert_called_once_with(rooms_trait._device_uid)
-        web_api_client.get_rooms.assert_not_called()
-    finally:
-        rooms_trait._home_data.rooms = original_rooms
+    assert home_data.received_devices
+    assert not rooms_trait.rooms
+
+    web_api_client.get_shared_device_rooms.return_value = [
+        HomeDataRoom(id=9999999, name="Office"),
+    ]
+    room_mapping_data = [[16, "2362048"], [17, "9999999"]]
+    mock_rpc_channel.send_command.side_effect = [room_mapping_data]
+
+    await rooms_trait.refresh()
+
+    assert rooms_trait.rooms == [
+        NamedRoomMapping(segment_id=16, iot_id="2362048", raw_name="Example room 1"),
+        NamedRoomMapping(segment_id=17, iot_id="9999999", raw_name="Office"),
+    ]
+    web_api_client.get_shared_device_rooms.assert_called_once_with(rooms_trait._device_uid)
+    web_api_client.get_rooms.assert_not_called()