Simplify roborock coordinator (#134700)

* Update roborock coordinator to require maps on startup

* Fix indent in merge

homeassistant/components/roborock/__init__.py

@@ -205,14 +205,6 @@ async def setup_device_v1(
     coordinator = RoborockDataUpdateCoordinator(
         hass, device, networking, product_info, mqtt_client, home_data_rooms
     )
-    # Verify we can communicate locally - if we can't, switch to cloud api
-    await coordinator.verify_api()
-    coordinator.api.is_available = True
-    try:
-        await coordinator.get_maps()
-    except RoborockException as err:
-        _LOGGER.warning("Failed to get map data")
-        _LOGGER.debug(err)
     try:
         await coordinator.async_config_entry_first_refresh()
     except ConfigEntryNotReady as ex:

homeassistant/components/roborock/coordinator.py

@@ -73,7 +73,27 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
         self.maps: dict[int, RoborockMapInfo] = {}
         self._home_data_rooms = {str(room.id): room.name for room in home_data_rooms}
 
-    async def verify_api(self) -> None:
+    async def _async_setup(self) -> None:
+        """Set up the coordinator."""
+        # Verify we can communicate locally - if we can't, switch to cloud api
+        await self._verify_api()
+        self.api.is_available = True
+
+        try:
+            maps = await self.api.get_multi_maps_list()
+        except RoborockException as err:
+            raise UpdateFailed("Failed to get map data: {err}") from err
+        # Rooms names populated later with calls to `set_current_map_rooms` for each map
+        self.maps = {
+            roborock_map.mapFlag: RoborockMapInfo(
+                flag=roborock_map.mapFlag,
+                name=roborock_map.name or f"Map {roborock_map.mapFlag}",
+                rooms={},
+            )
+            for roborock_map in (maps.map_info if (maps and maps.map_info) else ())
+        }
+
+    async def _verify_api(self) -> None:
         """Verify that the api is reachable. If it is not, switch clients."""
         if isinstance(self.api, RoborockLocalClientV1):
             try:
@@ -96,12 +116,8 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
 
     async def _update_device_prop(self) -> None:
         """Update device properties."""
-        device_prop = await self.api.get_prop()
-        if device_prop:
-            if self.roborock_device_info.props:
-                self.roborock_device_info.props.update(device_prop)
-            else:
-                self.roborock_device_info.props = device_prop
+        if (device_prop := await self.api.get_prop()) is not None:
+            self.roborock_device_info.props.update(device_prop)
 
     async def _async_update_data(self) -> DeviceProp:
         """Update data via library."""
@@ -111,7 +127,7 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
             # Set the new map id from the updated device props
             self._set_current_map()
             # Get the rooms for that map id.
-            await self.get_rooms()
+            await self.set_current_map_rooms()
         except RoborockException as ex:
             raise UpdateFailed(ex) from ex
         return self.roborock_device_info.props
@@ -127,29 +143,18 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
                 self.roborock_device_info.props.status.map_status - 3
             ) // 4
 
-    async def get_maps(self) -> None:
-        """Add a map to the coordinators mapping."""
-        maps = await self.api.get_multi_maps_list()
-        if maps and maps.map_info:
-            for roborock_map in maps.map_info:
-                self.maps[roborock_map.mapFlag] = RoborockMapInfo(
-                    flag=roborock_map.mapFlag,
-                    name=roborock_map.name or f"Map {roborock_map.mapFlag}",
-                    rooms={},
-                )
-
-    async def get_rooms(self) -> None:
-        """Get all of the rooms for the current map."""
+    async def set_current_map_rooms(self) -> None:
+        """Fetch all of the rooms for the current map and set on RoborockMapInfo."""
         # The api is only able to access rooms for the currently selected map
         # So it is important this is only called when you have the map you care
         # about selected.
-        if self.current_map in self.maps:
-            iot_rooms = await self.api.get_room_mapping()
-            if iot_rooms is not None:
-                for room in iot_rooms:
-                    self.maps[self.current_map].rooms[room.segment_id] = (
-                        self._home_data_rooms.get(room.iot_id, "Unknown")
-                    )
+        if self.current_map is None or self.current_map not in self.maps:
+            return
+        room_mapping = await self.api.get_room_mapping()
+        self.maps[self.current_map].rooms = {
+            room.segment_id: self._home_data_rooms.get(room.iot_id, "Unknown")
+            for room in room_mapping or ()
+        }
 
     @cached_property
     def duid(self) -> str:

homeassistant/components/roborock/image.py

@@ -121,7 +121,10 @@ class RoborockMap(RoborockCoordinatedEntityV1, ImageEntity):
         """Update the image if it is not cached."""
         if self.is_map_valid():
             response = await asyncio.gather(
-                *(self.cloud_api.get_map_v1(), self.coordinator.get_rooms()),
+                *(
+                    self.cloud_api.get_map_v1(),
+                    self.coordinator.set_current_map_rooms(),
+                ),
                 return_exceptions=True,
             )
             if not isinstance(response[0], bytes):
@@ -174,7 +177,8 @@ async def create_coordinator_maps(
             await asyncio.sleep(MAP_SLEEP)
         # Get the map data
         map_update = await asyncio.gather(
-            *[coord.cloud_api.get_map_v1(), coord.get_rooms()], return_exceptions=True
+            *[coord.cloud_api.get_map_v1(), coord.set_current_map_rooms()],
+            return_exceptions=True,
         )
         # If we fail to get the map, we should set it to empty byte,
         # still create it, and set it as unavailable.

tests/components/roborock/test_init.py

@@ -133,20 +133,18 @@ async def test_local_client_fails_props(
         assert mock_roborock_entry.state is ConfigEntryState.SETUP_RETRY
 
 
-async def test_fails_maps_continue(
+async def test_fail_maps(
     hass: HomeAssistant,
     mock_roborock_entry: MockConfigEntry,
     bypass_api_fixture_v1_only,
 ) -> None:
-    """Test that if we fail to get the maps, we still setup."""
+    """Test that the integration fails to load if we fail to get the maps."""
     with patch(
         "homeassistant.components.roborock.coordinator.RoborockLocalClientV1.get_multi_maps_list",
         side_effect=RoborockException(),
     ):
         await async_setup_component(hass, DOMAIN, {})
-        assert mock_roborock_entry.state is ConfigEntryState.LOADED
-        # No map data means no images
-        assert len(hass.states.async_all("image")) == 0
+        assert mock_roborock_entry.state is ConfigEntryState.SETUP_RETRY
 
 
 async def test_reauth_started(