Move Roborock map refresh to coordinator (#140758)

Move refresh coordinator to coordinator

homeassistant/components/roborock/__init__.py

@@ -111,6 +111,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) ->
             translation_key="no_coordinators",
         )
     valid_coordinators = RoborockCoordinators(v1_coords, a01_coords)
+    await asyncio.gather(
+        *(coord.refresh_coordinator_map() for coord in valid_coordinators.v1)
+    )
 
     async def on_stop(_: Any) -> None:
         _LOGGER.debug("Shutting down roborock")

homeassistant/components/roborock/coordinator.py

@@ -48,6 +48,7 @@ from .const import (
     DRAWABLES,
     MAP_FILE_FORMAT,
     MAP_SCALE,
+    MAP_SLEEP,
     V1_CLOUD_IN_CLEANING_INTERVAL,
     V1_CLOUD_NOT_CLEANING_INTERVAL,
     V1_LOCAL_IN_CLEANING_INTERVAL,
@@ -316,6 +317,36 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
         """Get the slug of the duid."""
         return slugify(self.duid)
 
+    async def refresh_coordinator_map(self) -> None:
+        """Get the starting map information for all maps for this device.
+
+        The following steps must be done synchronously.
+        Only one map can be loaded at a time per device.
+        """
+        cur_map = self.current_map
+        # This won't be None at this point as the coordinator will have run first.
+        if cur_map is None:
+            # If we don't have a cur map(shouldn't happen) just
+            # return as we can't do anything.
+            return
+        map_flags = sorted(self.maps, key=lambda data: data == cur_map, reverse=True)
+        for map_flag in map_flags:
+            if map_flag != cur_map:
+                # Only change the map and sleep if we have multiple maps.
+                await self.api.load_multi_map(map_flag)
+                self.current_map = map_flag
+                # We cannot get the map until the roborock servers fully process the
+                # map change.
+                await asyncio.sleep(MAP_SLEEP)
+            await self.set_current_map_rooms()
+
+        if len(self.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 self.api.load_multi_map(cur_map)
+            self.current_map = cur_map
+
 
 class RoborockDataUpdateCoordinatorA01(
     DataUpdateCoordinator[

homeassistant/components/roborock/image.py

@@ -4,8 +4,6 @@ import asyncio
 from datetime import datetime
 import logging
 
-from roborock import RoborockCommand
-
 from homeassistant.components.image import ImageEntity
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
@@ -14,7 +12,7 @@ from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.util import dt as dt_util
 
-from .const import DOMAIN, IMAGE_CACHE_INTERVAL, MAP_SLEEP
+from .const import DOMAIN, IMAGE_CACHE_INTERVAL
 from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
 from .entity import RoborockCoordinatedEntityV1
 
@@ -28,9 +26,6 @@ async def async_setup_entry(
 ) -> None:
     """Set up Roborock image platform."""
 
-    await asyncio.gather(
-        *(refresh_coordinators(hass, coord) for coord in config_entry.runtime_data.v1)
-    )
     async_add_entities(
         (
             RoborockMap(
@@ -126,33 +121,3 @@ class RoborockMap(RoborockCoordinatedEntityV1, ImageEntity):
                     content,
                 )
         return self.cached_map
-
-
-async def refresh_coordinators(
-    hass: HomeAssistant, coord: RoborockDataUpdateCoordinator
-) -> None:
-    """Get the starting map information for all maps for this device.
-
-    The following steps must be done synchronously.
-    Only one map can be loaded at a time per device.
-    """
-    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
-    map_flags = sorted(coord.maps, key=lambda data: data == cur_map, reverse=True)
-    for map_flag in map_flags:
-        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])
-            coord.current_map = map_flag
-            # We cannot get the map until the roborock servers fully process the
-            # map change.
-            await asyncio.sleep(MAP_SLEEP)
-        await coord.set_current_map_rooms()
-
-    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])
-        coord.current_map = cur_map

tests/components/roborock/conftest.py

@@ -80,6 +80,9 @@ def bypass_api_client_fixture() -> None:
             "homeassistant.components.roborock.RoborockApiClient.get_scenes",
             return_value=SCENES,
         ),
+        patch(
+            "homeassistant.components.roborock.coordinator.RoborockLocalClientV1.load_multi_map"
+        ),
     ):
         yield
 
@@ -127,7 +130,7 @@ def bypass_api_fixture(bypass_api_client_fixture: Any) -> None:
             "roborock.version_1_apis.AttributeCache.value",
         ),
         patch(
-            "homeassistant.components.roborock.image.MAP_SLEEP",
+            "homeassistant.components.roborock.coordinator.MAP_SLEEP",
             0,
         ),
         patch(

tests/components/roborock/test_image.py

@@ -244,3 +244,33 @@ async def test_fail_updating_image(
         async_fire_time_changed(hass, now)
         resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
     assert not resp.ok
+
+
+async def test_index_error_map(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    hass_client: ClientSessionGenerator,
+) -> None:
+    """Test that we handle failing getting the image after it has already been setup with a indexerror."""
+    client = await hass_client()
+    now = dt_util.utcnow() + timedelta(seconds=91)
+    # Copy the device prop so we don't override it
+    prop = copy.deepcopy(PROP)
+    prop.status.in_cleaning = 1
+    # Update image, but get IndexError for image.
+    with (
+        patch(
+            "homeassistant.components.roborock.coordinator.RoborockMapDataParser.parse",
+            side_effect=IndexError,
+        ),
+        patch(
+            "homeassistant.components.roborock.coordinator.RoborockLocalClientV1.get_prop",
+            return_value=prop,
+        ),
+        patch(
+            "homeassistant.components.roborock.image.dt_util.utcnow", return_value=now
+        ),
+    ):
+        async_fire_time_changed(hass, now)
+        resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
+    assert not resp.ok