Move Roborock map retrieval to coordinator and made map always diagnostic (#104680)

Co-authored-by: Robert Resch <robert@resch.dev>

homeassistant/components/roborock/__init__.py

@@ -122,6 +122,11 @@ async def setup_device(
     # 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

@@ -59,6 +59,8 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
 
         if mac := self.roborock_device_info.network_info.mac:
             self.device_info[ATTR_CONNECTIONS] = {(dr.CONNECTION_NETWORK_MAC, mac)}
+        # Maps from map flag to map name
+        self.maps: dict[int, str] = {}
 
     async def verify_api(self) -> None:
         """Verify that the api is reachable. If it is not, switch clients."""
@@ -107,3 +109,10 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
             self.current_map = (
                 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] = roborock_map.name

homeassistant/components/roborock/image.py

@@ -66,13 +66,7 @@ class RoborockMap(RoborockCoordinatedEntity, ImageEntity):
         self._attr_image_last_updated = dt_util.utcnow()
         self.map_flag = map_flag
         self.cached_map = self._create_image(starting_map)
-
-    @property
-    def entity_category(self) -> EntityCategory | None:
-        """Return diagnostic entity category for any non-selected maps."""
-        if not self.is_selected:
-            return EntityCategory.DIAGNOSTIC
-        return None
+        self._attr_entity_category = EntityCategory.DIAGNOSTIC
 
     @property
     def is_selected(self) -> bool:
@@ -127,42 +121,37 @@ async def create_coordinator_maps(
     Only one map can be loaded at a time per device.
     """
     entities = []
-    maps = await coord.cloud_api.get_multi_maps_list()
-    if maps is not None and maps.map_info is not None:
-        cur_map = coord.current_map
-        # This won't be None at this point as the coordinator will have run first.
-        assert cur_map is not None
-        # Sort the maps so that we start with the current map and we can skip the
-        # load_multi_map call.
-        maps_info = sorted(
-            maps.map_info, key=lambda data: data.mapFlag == cur_map, reverse=True
+
+    cur_map = coord.current_map
+    # This won't be None at this point as the coordinator will have run first.
+    assert cur_map is not None
+    # Sort the maps so that we start with the current map and we can skip the
+    # load_multi_map call.
+    maps_info = sorted(
+        coord.maps.items(), key=lambda data: data[0] == cur_map, reverse=True
+    )
+    for map_flag, map_name in maps_info:
+        # Load the map - so we can access it with get_map_v1
+        if map_flag != cur_map:
+            # Only change the map and sleep if we have multiple maps.
+            await coord.api.send_command(RoborockCommand.LOAD_MULTI_MAP, [map_flag])
+            # We cannot get the map until the roborock servers fully process the
+            # map change.
+            await asyncio.sleep(MAP_SLEEP)
+        # Get the map data
+        api_data: bytes = await coord.cloud_api.get_map_v1()
+        entities.append(
+            RoborockMap(
+                f"{slugify(coord.roborock_device_info.device.duid)}_map_{map_name}",
+                coord,
+                map_flag,
+                api_data,
+                map_name,
+            )
         )
-        for roborock_map in maps_info:
-            # Load the map - so we can access it with get_map_v1
-            if roborock_map.mapFlag != cur_map:
-                # Only change the map and sleep if we have multiple maps.
-                await coord.api.send_command(
-                    RoborockCommand.LOAD_MULTI_MAP, [roborock_map.mapFlag]
-                )
-                # We cannot get the map until the roborock servers fully process the
-                # map change.
-                await asyncio.sleep(MAP_SLEEP)
-            # Get the map data
-            api_data: bytes = await coord.cloud_api.get_map_v1()
-            entities.append(
-                RoborockMap(
-                    f"{slugify(coord.roborock_device_info.device.duid)}_map_{roborock_map.name}",
-                    coord,
-                    roborock_map.mapFlag,
-                    api_data,
-                    roborock_map.name,
-                )
-            )
-        if len(maps.map_info) != 1:
-            # Set the map back to the map the user previously had selected so that it
-            # does not change the end user's app.
-            # Only needs to happen when we changed maps above.
-            await coord.cloud_api.send_command(
-                RoborockCommand.LOAD_MULTI_MAP, [cur_map]
-            )
+    if len(coord.maps) != 1:
+        # Set the map back to the map the user previously had selected so that it
+        # does not change the end user's app.
+        # Only needs to happen when we changed maps above.
+        await coord.cloud_api.send_command(RoborockCommand.LOAD_MULTI_MAP, [cur_map])
     return entities

tests/components/roborock/conftest.py

@@ -45,6 +45,9 @@ def bypass_api_fixture() -> None:
     ), patch(
         "homeassistant.components.roborock.coordinator.RoborockMqttClient.get_multi_maps_list",
         return_value=MULTI_MAP_LIST,
+    ), patch(
+        "homeassistant.components.roborock.coordinator.RoborockLocalClient.get_multi_maps_list",
+        return_value=MULTI_MAP_LIST,
     ), patch(
         "homeassistant.components.roborock.image.RoborockMapDataParser.parse",
         return_value=MAP_DATA,

tests/components/roborock/test_init.py

@@ -107,6 +107,20 @@ async def test_local_client_fails_props(
         assert mock_roborock_entry.state is ConfigEntryState.SETUP_RETRY
 
 
+async def test_fails_maps_continue(
+    hass: HomeAssistant, mock_roborock_entry: MockConfigEntry, bypass_api_fixture
+) -> None:
+    """Test that if we fail to get the maps, we still setup."""
+    with patch(
+        "homeassistant.components.roborock.coordinator.RoborockLocalClient.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
+
+
 async def test_reauth_started(
     hass: HomeAssistant, bypass_api_fixture, mock_roborock_entry: MockConfigEntry
 ) -> None: