chore: address comments

Author: Lash-L

roborock/devices/traits/v1/__init__.py

@@ -193,10 +193,10 @@ def __init__(
         self.device_features = DeviceFeaturesTrait(product, self._device_cache)
         self.status = StatusTrait(self.device_features, region=self._region)
         self.consumables = ConsumableTrait()
-        self.rooms = RoomsTrait(home_data)
+        self.rooms = RoomsTrait(home_data, web_api)
         self.maps = MapsTrait(self.status)
         self.map_content = MapContentTrait(map_parser_config)
-        self.home = HomeTrait(self.status, self.maps, self.map_content, self.rooms, self._device_cache, web_api)
+        self.home = HomeTrait(self.status, self.maps, self.map_content, self.rooms, self._device_cache)
         self.network_info = NetworkInfoTrait(device_uid, self._device_cache)
         self.routines = RoutinesTrait(device_uid, web_api)
 

roborock/devices/traits/v1/home.py

@@ -26,7 +26,6 @@
 from roborock.devices.traits.v1 import common
 from roborock.exceptions import RoborockDeviceBusy, RoborockException, RoborockInvalidStatus
 from roborock.roborock_typing import RoborockCommand
-from roborock.web_api import UserWebApiClient
 
 from .map_content import MapContent, MapContentTrait
 from .maps import MapsTrait
@@ -50,7 +49,6 @@ def __init__(
         map_content: MapContentTrait,
         rooms_trait: RoomsTrait,
         device_cache: DeviceCache,
-        web_api: UserWebApiClient | None = None,
     ) -> None:
         """Initialize the HomeTrait.
 
@@ -73,8 +71,6 @@ def __init__(
         self._map_content = map_content
         self._rooms_trait = rooms_trait
         self._device_cache = device_cache
-        self._web_api = web_api
-        self._seen_room_iot_ids_by_map: dict[int, set[str]] = {}
         self._discovery_completed = False
         self._home_map_info: dict[int, CombinedMapInfo] | None = None
         self._home_map_content: dict[int, MapContent] | None = None
@@ -94,7 +90,6 @@ async def discover_home(self) -> None:
         if device_cache_data and device_cache_data.home_map_info:
             _LOGGER.debug("Home cache already populated, skipping discovery")
             self._home_map_info = device_cache_data.home_map_info
-            self._seed_seen_room_iot_ids(device_cache_data.home_map_info)
             self._discovery_completed = True
             try:
                 self._home_map_content = {
@@ -143,39 +138,7 @@ async def _refresh_map_info(self, map_info) -> CombinedMapInfo:
                         # or if the room name isn't unknown.
                         rooms[room.segment_id] = room
 
-        # If we discovered at least one new unknown room iot_id for this map, fetch names from web API.
-        seen_room_iot_ids = self._seen_room_iot_ids_by_map.setdefault(map_info.map_flag, set())
-        has_new_unknown_room_iot_id = any(
-            room.name == "Unknown" and room.iot_id not in seen_room_iot_ids for room in rooms.values()
-        )
-        if self._web_api and has_new_unknown_room_iot_id:
-            try:
-                web_rooms = await self._web_api.get_rooms()
-                web_room_names = {str(r.id): r.name for r in web_rooms}
-                for segment_id, room in rooms.items():
-                    if room.name == "Unknown" and room.iot_id in web_room_names:
-                        rooms[segment_id] = NamedRoomMapping(
-                            segment_id=room.segment_id,
-                            iot_id=room.iot_id,
-                            name=web_room_names[room.iot_id],
-                        )
-                # Merge new rooms into home_data so future refreshes benefit
-                if web_rooms:
-                    self._rooms_trait.merge_home_data_rooms(web_rooms)
-            except Exception:
-                # Broad exception as we don't want anything here to make us fail upwards, we are okay with 'unknowns'
-                _LOGGER.debug("Failed to fetch rooms from web API for map %s", map_info.map_flag)
-
-        # Replace remaining "Unknown" names with "Room {segment_id}" fallback.
-        for segment_id, room in rooms.items():
-            if room.name == "Unknown":
-                rooms[segment_id] = NamedRoomMapping(
-                    segment_id=room.segment_id,
-                    iot_id=room.iot_id,
-                    name=f"Room {room.segment_id}",
-                )
-
-        self._seen_room_iot_ids_by_map[map_info.map_flag].update(room.iot_id for room in rooms.values())
+        await self._rooms_trait.resolve_unknown_room_names(rooms)
 
         return CombinedMapInfo(
             map_flag=map_info.map_flag,
@@ -293,7 +256,6 @@ async def _update_home_cache(
         await self._device_cache.set(device_cache_data)
         self._home_map_info = home_map_info
         self._home_map_content = home_map_content
-        self._seed_seen_room_iot_ids(home_map_info)
 
     async def _update_current_map(
         self,
@@ -322,13 +284,7 @@ async def _update_current_map(
         if self._home_map_info is None:
             self._home_map_info = {}
         self._home_map_info[map_flag] = map_info
-        self._seed_seen_room_iot_ids({map_flag: map_info})
 
         if self._home_map_content is None:
             self._home_map_content = {}
         self._home_map_content[map_flag] = map_content
-
-    def _seed_seen_room_iot_ids(self, home_map_info: dict[int, CombinedMapInfo]) -> None:
-        """Seed known room iot_ids from cached/current map information."""
-        for map_flag, map_info in home_map_info.items():
-            self._seen_room_iot_ids_by_map.setdefault(map_flag, set()).update(room.iot_id for room in map_info.rooms)

roborock/devices/traits/v1/rooms.py

@@ -6,6 +6,7 @@
 from roborock.data import HomeData, HomeDataRoom, NamedRoomMapping, RoborockBase
 from roborock.devices.traits.v1 import common
 from roborock.roborock_typing import RoborockCommand
+from roborock.web_api import UserWebApiClient
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -32,10 +33,12 @@ class RoomsTrait(Rooms, common.V1TraitMixin):
 
     command = RoborockCommand.GET_ROOM_MAPPING
 
-    def __init__(self, home_data: HomeData) -> None:
+    def __init__(self, home_data: HomeData, web_api: UserWebApiClient | None = None) -> None:
         """Initialize the RoomsTrait."""
         super().__init__()
         self._home_data = home_data
+        self._web_api = web_api
+        self._seen_unknown_room_iot_ids: set[str] = set()
 
     @property
     def _iot_id_room_name_map(self) -> dict[str, str]:
@@ -57,14 +60,46 @@ def _parse_response(self, response: common.V1ResponseData) -> Rooms:
 
     def merge_home_data_rooms(self, rooms: list[HomeDataRoom]) -> None:
         """Merge newly discovered rooms into home data by room id."""
-        existing_ids = {room.id for room in self._home_data.rooms or ()}
         updated_rooms = list(self._home_data.rooms or ())
+        existing_by_id = {room.id: room for room in updated_rooms}
+
         for room in rooms:
-            if room.id not in existing_ids:
+            existing_room = existing_by_id.get(room.id)
+            if existing_room is None:
                 updated_rooms.append(room)
-                existing_ids.add(room.id)
+                existing_by_id[room.id] = room
+            elif room.name and existing_room.name in ("", _DEFAULT_NAME):
+                existing_room.name = room.name
+
         self._home_data.rooms = updated_rooms
 
+    async def resolve_unknown_room_names(self, rooms: dict[int, NamedRoomMapping]) -> None:
+        """Resolve unknown room names using home data and web API fallbacks."""
+        unknown_room_iot_ids = {room.iot_id for room in rooms.values() if room.name == _DEFAULT_NAME}
+        new_unknown_room_iot_ids = unknown_room_iot_ids - self._seen_unknown_room_iot_ids
+        web_room_names: dict[str, str] = {}
+
+        if self._web_api and new_unknown_room_iot_ids:
+            try:
+                web_rooms = await self._web_api.get_rooms()
+            except Exception:
+                _LOGGER.debug("Failed to fetch rooms from web API", exc_info=True)
+            else:
+                if web_rooms:
+                    web_room_names = {str(room.id): room.name for room in web_rooms}
+                    self.merge_home_data_rooms(web_rooms)
+
+        for segment_id, room in rooms.items():
+            if room.name != _DEFAULT_NAME:
+                continue
+            rooms[segment_id] = NamedRoomMapping(
+                segment_id=room.segment_id,
+                iot_id=room.iot_id,
+                name=web_room_names.get(room.iot_id, f"Room {room.segment_id}"),
+            )
+
+        self._seen_unknown_room_iot_ids.update(unknown_room_iot_ids)
+
 
 def _extract_segment_pairs(response: list) -> list[tuple[int, str]]:
     """Extract segment_id and iot_id pairs from the response.

tests/devices/traits/v1/test_home.py

@@ -157,7 +157,8 @@ def home_trait(
     mock_web_api: AsyncMock,
 ) -> HomeTrait:
     """Create a HomeTrait instance with mocked dependencies."""
-    return HomeTrait(status_trait, maps_trait, map_content_trait, rooms_trait, device_cache, mock_web_api)
+    rooms_trait._web_api = mock_web_api
+    return HomeTrait(status_trait, maps_trait, map_content_trait, rooms_trait, device_cache)
 
 
 @pytest.fixture(autouse=True)

tests/devices/traits/v1/test_rooms.py

@@ -5,6 +5,7 @@
 
 import pytest
 
+from roborock.data.containers import HomeDataRoom, NamedRoomMapping
 from roborock.devices.device import RoborockDevice
 from roborock.devices.traits.v1.rooms import RoomsTrait
 from roborock.devices.traits.v1.status import StatusTrait
@@ -70,3 +71,56 @@ async def test_refresh_rooms_trait(
     # 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)
+
+
+def test_merge_home_data_rooms_appends_new_rooms(rooms_trait: RoomsTrait) -> None:
+    """Test merge_home_data_rooms appends new rooms without replacing known names."""
+    rooms_trait.merge_home_data_rooms(
+        [
+            HomeDataRoom(id=2362048, name="Living Room"),
+            HomeDataRoom(id=9999999, name="Office"),
+        ]
+    )
+
+    home_data_rooms = {str(room.id): room.name for room in rooms_trait._home_data.rooms}
+    assert home_data_rooms["2362048"] == "Example room 1"
+    assert home_data_rooms["9999999"] == "Office"
+
+
+async def test_resolve_unknown_room_names_web_api_called_once(
+    rooms_trait: RoomsTrait,
+    web_api_client: AsyncMock,
+) -> None:
+    """Test unknown room IDs trigger one web lookup per iot_id."""
+    web_api_client.get_rooms.return_value = [
+        HomeDataRoom(id=2362048, name="Living Room"),
+    ]
+
+    room_map = {
+        16: NamedRoomMapping(segment_id=16, iot_id="2362048", name="Unknown"),
+    }
+    await rooms_trait.resolve_unknown_room_names(room_map)
+    assert room_map[16].name == "Living Room"
+
+    second_room_map = {
+        16: NamedRoomMapping(segment_id=16, iot_id="2362048", name="Unknown"),
+    }
+    await rooms_trait.resolve_unknown_room_names(second_room_map)
+    assert second_room_map[16].name == "Room 16"
+    web_api_client.get_rooms.assert_called_once()
+
+
+async def test_resolve_unknown_room_names_falls_back_to_segment_id(
+    rooms_trait: RoomsTrait,
+    web_api_client: AsyncMock,
+) -> None:
+    """Test unresolved unknown names use Room {segment_id} fallback."""
+    web_api_client.get_rooms.return_value = []
+    room_map = {
+        33: NamedRoomMapping(segment_id=33, iot_id="9999911", name="Unknown"),
+    }
+
+    await rooms_trait.resolve_unknown_room_names(room_map)
+
+    assert room_map[33].name == "Room 33"
+    web_api_client.get_rooms.assert_called_once()