Move MapData to Coordinator for Roborock (#140766)

* Move MapData to Coordinator

* seeing if mypy likes this

* delete dead code

* Some MR comments

* remove MapData and always update on startup if we don't have a stored map.

* don't do on demand updates

* remove unneeded logic and pull out map save

* Apply suggestions from code review

Co-authored-by: Allen Porter <allen.porter@gmail.com>

* see if mypy is happy

---------

Co-authored-by: Allen Porter <allen.porter@gmail.com>

homeassistant/components/roborock/const.py

@@ -46,7 +46,7 @@ PLATFORMS = [
 ]
 
 # This can be lowered in the future if we do not receive rate limiting issues.
-IMAGE_CACHE_INTERVAL = 30
+IMAGE_CACHE_INTERVAL = timedelta(seconds=30)
 
 MAP_SLEEP = 3
 

homeassistant/components/roborock/coordinator.py

@@ -39,13 +39,14 @@ from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
-from homeassistant.util import slugify
+from homeassistant.util import dt as dt_util, slugify
 
 from .const import (
     A01_UPDATE_INTERVAL,
     DEFAULT_DRAWABLES,
     DOMAIN,
     DRAWABLES,
+    IMAGE_CACHE_INTERVAL,
     MAP_FILE_FORMAT,
     MAP_SCALE,
     MAP_SLEEP,
@@ -191,15 +192,59 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
         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
+        roborock_maps = maps.map_info if (maps and maps.map_info) else ()
+        stored_images = await asyncio.gather(
+            *[
+                self.map_storage.async_load_map(roborock_map.mapFlag)
+                for roborock_map in roborock_maps
+            ]
+        )
         self.maps = {
             roborock_map.mapFlag: RoborockMapInfo(
                 flag=roborock_map.mapFlag,
                 name=roborock_map.name or f"Map {roborock_map.mapFlag}",
                 rooms={},
+                image=image,
+                last_updated=dt_util.utcnow() - IMAGE_CACHE_INTERVAL,
             )
-            for roborock_map in (maps.map_info if (maps and maps.map_info) else ())
+            for image, roborock_map in zip(stored_images, roborock_maps, strict=False)
         }
 
+    async def update_map(self) -> None:
+        """Update the currently selected map."""
+        # The current map was set in the props update, so these can be done without
+        # worry of applying them to the wrong map.
+        if self.current_map is None:
+            # This exists as a safeguard/ to keep mypy happy.
+            return
+        try:
+            response = await self.cloud_api.get_map_v1()
+        except RoborockException as ex:
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="map_failure",
+            ) from ex
+        if not isinstance(response, bytes):
+            _LOGGER.debug("Failed to parse map contents: %s", response)
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="map_failure",
+            )
+        parsed_image = self.parse_image(response)
+        if parsed_image is None:
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="map_failure",
+            )
+        if parsed_image != self.maps[self.current_map].image:
+            await self.map_storage.async_save_map(
+                self.current_map,
+                parsed_image,
+            )
+            current_roborock_map_info = self.maps[self.current_map]
+            current_roborock_map_info.image = parsed_image
+            current_roborock_map_info.last_updated = dt_util.utcnow()
+
     async def _verify_api(self) -> None:
         """Verify that the api is reachable. If it is not, switch clients."""
         if isinstance(self.api, RoborockLocalClientV1):
@@ -240,6 +285,19 @@ 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.
+
+            # If the vacuum is currently cleaning and it has been IMAGE_CACHE_INTERVAL
+            # since the last map update, you can update the map.
+            if (
+                self.current_map is not None
+                and self.roborock_device_info.props.status.in_cleaning
+                and (dt_util.utcnow() - self.maps[self.current_map].last_updated)
+                > IMAGE_CACHE_INTERVAL
+            ):
+                try:
+                    await self.update_map()
+                except HomeAssistantError as err:
+                    _LOGGER.debug("Failed to update map: %s", err)
             await self.set_current_map_rooms()
         except RoborockException as ex:
             _LOGGER.debug("Failed to update data: %s", ex)
@@ -338,7 +396,14 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
                 # 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()
+            tasks = [self.set_current_map_rooms()]
+            # The image is set within async_setup, so if it exists, we have it here.
+            if self.maps[map_flag].image is None:
+                # If we don't have a cached map, let's update it here so that it can be
+                # cached in the future.
+                tasks.append(self.update_map())
+            # If either of these fail, we don't care, and we want to continue.
+            await asyncio.gather(*tasks, return_exceptions=True)
 
         if len(self.maps) != 1:
             # Set the map back to the map the user previously had selected so that it

homeassistant/components/roborock/image.py

@@ -1,6 +1,5 @@
 """Support for Roborock image."""
 
-import asyncio
 from datetime import datetime
 import logging
 
@@ -8,11 +7,8 @@ from homeassistant.components.image import ImageEntity
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
-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
 from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
 from .entity import RoborockCoordinatedEntityV1
 
@@ -75,51 +71,19 @@ class RoborockMap(RoborockCoordinatedEntityV1, ImageEntity):
     async def async_added_to_hass(self) -> None:
         """When entity is added to hass load any previously cached maps from disk."""
         await super().async_added_to_hass()
-        content = await self.coordinator.map_storage.async_load_map(self.map_flag)
+        self._attr_image_last_updated = self.coordinator.maps[
-        self.cached_map = content or b""
+            self.map_flag
-        self._attr_image_last_updated = dt_util.utcnow()
+        ].last_updated
         self.async_write_ha_state()
 
     def _handle_coordinator_update(self) -> None:
-        # Bump last updated every third time the coordinator runs, so that async_image
+        # If the coordinator has updated the map, we can update the image.
-        # will be called and we will evaluate on the new coordinator data if we should
+        self._attr_image_last_updated = self.coordinator.maps[
-        # update the cache.
+            self.map_flag
-        if self.is_selected and (
+        ].last_updated
-            (
+
-                (dt_util.utcnow() - self.image_last_updated).total_seconds()
-                > IMAGE_CACHE_INTERVAL
-                and self.coordinator.roborock_device_info.props.status is not None
-                and bool(self.coordinator.roborock_device_info.props.status.in_cleaning)
-            )
-            or self.cached_map == b""
-        ):
-            # This will tell async_image it should update.
-            self._attr_image_last_updated = dt_util.utcnow()
         super()._handle_coordinator_update()
 
     async def async_image(self) -> bytes | None:
-        """Update the image if it is not cached."""
+        """Get the cached image."""
-        if self.is_selected:
+        return self.coordinator.maps[self.map_flag].image
-            response = await asyncio.gather(
-                *(
-                    self.cloud_api.get_map_v1(),
-                    self.coordinator.set_current_map_rooms(),
-                ),
-                return_exceptions=True,
-            )
-            if (
-                not isinstance(response[0], bytes)
-                or (content := self.coordinator.parse_image(response[0])) is None
-            ):
-                _LOGGER.debug("Failed to parse map contents: %s", response[0])
-                raise HomeAssistantError(
-                    translation_domain=DOMAIN,
-                    translation_key="map_failure",
-                )
-            if self.cached_map != content:
-                self.cached_map = content
-                await self.coordinator.map_storage.async_save_map(
-                    self.map_flag,
-                    content,
-                )
-        return self.cached_map

homeassistant/components/roborock/models.py

@@ -1,6 +1,7 @@
 """Roborock Models."""
 
 from dataclasses import dataclass
+from datetime import datetime
 from typing import Any
 
 from roborock.containers import HomeDataDevice, HomeDataProduct, NetworkInfo
@@ -48,3 +49,5 @@ class RoborockMapInfo:
     flag: int
     name: str
     rooms: dict[int, str]
+    image: bytes | None
+    last_updated: datetime

homeassistant/components/roborock/vacuum.py

@@ -1,6 +1,5 @@
 """Support for Roborock vacuum class."""
 
-from dataclasses import asdict
 from typing import Any
 
 from roborock.code_mappings import RoborockStateCode
@@ -206,7 +205,14 @@ class RoborockVacuum(RoborockCoordinatedEntityV1, StateVacuumEntity):
         """Get map information such as map id and room ids."""
         return {
             "maps": [
-                asdict(vacuum_map) for vacuum_map in self.coordinator.maps.values()
+                {
+                    "flag": vacuum_map.flag,
+                    "name": vacuum_map.name,
+                    # JsonValueType does not accept a int as a key - was not a
+                    # issue with previous asdict() implementation.
+                    "rooms": vacuum_map.rooms,  # type: ignore[dict-item]
+                }
+                for vacuum_map in self.coordinator.maps.values()
             ]
         }
 

tests/components/roborock/conftest.py

@@ -228,7 +228,7 @@ async def setup_entry(
         yield mock_roborock_entry
 
 
-@pytest.fixture
+@pytest.fixture(autouse=True)
 async def cleanup_map_storage(
     hass: HomeAssistant, mock_roborock_entry: MockConfigEntry
 ) -> Generator[pathlib.Path]:

tests/components/roborock/test_config_flow.py

@@ -25,6 +25,12 @@ from .mock_data import MOCK_CONFIG, USER_DATA, USER_EMAIL
 from tests.common import MockConfigEntry
 
 
+@pytest.fixture
+def cleanup_map_storage():
+    """Override the map storage fixture as it is not relevant here."""
+    return
+
+
 async def test_config_flow_success(
     hass: HomeAssistant,
     bypass_api_fixture,
@@ -189,25 +195,31 @@ async def test_config_flow_failures_code_login(
 
 
 async def test_options_flow_drawables(
-    hass: HomeAssistant, setup_entry: MockConfigEntry
+    hass: HomeAssistant, mock_roborock_entry: MockConfigEntry
 ) -> None:
     """Test that the options flow works."""
-    result = await hass.config_entries.options.async_init(setup_entry.entry_id)
+    with patch("homeassistant.components.roborock.roborock_storage"):
-
+        await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
-    assert result["type"] == FlowResultType.FORM
-    assert result["step_id"] == DRAWABLES
-    with patch(
-        "homeassistant.components.roborock.async_setup_entry", return_value=True
-    ) as mock_setup:
-        result = await hass.config_entries.options.async_configure(
-            result["flow_id"],
-            user_input={Drawable.PREDICTED_PATH: True},
-        )
         await hass.async_block_till_done()
 
-    assert result["type"] == FlowResultType.CREATE_ENTRY
+        result = await hass.config_entries.options.async_init(
-    assert setup_entry.options[DRAWABLES][Drawable.PREDICTED_PATH] is True
+            mock_roborock_entry.entry_id
-    assert len(mock_setup.mock_calls) == 1
+        )
+
+        assert result["type"] == FlowResultType.FORM
+        assert result["step_id"] == DRAWABLES
+        with patch(
+            "homeassistant.components.roborock.async_setup_entry", return_value=True
+        ) as mock_setup:
+            result = await hass.config_entries.options.async_configure(
+                result["flow_id"],
+                user_input={Drawable.PREDICTED_PATH: True},
+            )
+            await hass.async_block_till_done()
+
+        assert result["type"] == FlowResultType.CREATE_ENTRY
+        assert mock_roborock_entry.options[DRAWABLES][Drawable.PREDICTED_PATH] is True
+        assert len(mock_setup.mock_calls) == 1
 
 
 async def test_reauth_flow(

tests/components/roborock/test_image.py

@@ -62,20 +62,26 @@ async def test_floorplan_image(
             return_value=prop,
         ),
         patch(
-            "homeassistant.components.roborock.image.dt_util.utcnow", return_value=now
+            "homeassistant.components.roborock.coordinator.dt_util.utcnow",
+            return_value=now,
         ),
         patch(
             "homeassistant.components.roborock.coordinator.RoborockMapDataParser.parse",
-            return_value=new_map_data,
+            return_value=MAP_DATA,
         ) as parse_map,
     ):
+        # This should call parse_map twice as the both devices are in cleaning.
         async_fire_time_changed(hass, now)
-        await hass.async_block_till_done()
         resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
+        assert resp.status == HTTPStatus.OK
+        resp = await client.get("/api/image_proxy/image.roborock_s7_2_upstairs")
+        assert resp.status == HTTPStatus.OK
+        resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_downstairs")
     assert resp.status == HTTPStatus.OK
     body = await resp.read()
     assert body is not None
-    assert parse_map.call_count == 1
+
+    assert parse_map.call_count == 2
 
 
 async def test_floorplan_image_failed_parse(
@@ -91,6 +97,7 @@ async def test_floorplan_image_failed_parse(
     # Copy the device prop so we don't override it
     prop = copy.deepcopy(PROP)
     prop.status.in_cleaning = 1
+    previous_state = hass.states.get("image.roborock_s7_maxv_upstairs").state
     # Update image, but get none for parse image.
     with (
         patch(
@@ -102,12 +109,16 @@ async def test_floorplan_image_failed_parse(
             return_value=prop,
         ),
         patch(
-            "homeassistant.components.roborock.image.dt_util.utcnow", return_value=now
+            "homeassistant.components.roborock.coordinator.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
+    # The map should load fine from the coordinator, but it should not update the
+    # last_updated timestamp.
+    assert resp.ok
+    assert previous_state == hass.states.get("image.roborock_s7_maxv_upstairs").state
 
 
 async def test_fail_to_save_image(
@@ -158,6 +169,9 @@ async def test_fail_to_load_image(
             "homeassistant.components.roborock.roborock_storage.Path.read_bytes",
             side_effect=OSError,
         ) as read_bytes,
+        patch(
+            "homeassistant.components.roborock.coordinator.RoborockDataUpdateCoordinator.refresh_coordinator_map"
+        ),
     ):
         # Reload the config entry so that the map is saved in storage and entities exist.
         await hass.config_entries.async_reload(setup_entry.entry_id)
@@ -224,6 +238,7 @@ async def test_fail_updating_image(
     prop = copy.deepcopy(PROP)
     prop.status.in_cleaning = 1
     # Update image, but get none for parse image.
+    previous_state = hass.states.get("image.roborock_s7_maxv_upstairs").state
     with (
         patch(
             "homeassistant.components.roborock.coordinator.RoborockMapDataParser.parse",
@@ -234,7 +249,8 @@ async def test_fail_updating_image(
             return_value=prop,
         ),
         patch(
-            "homeassistant.components.roborock.image.dt_util.utcnow", return_value=now
+            "homeassistant.components.roborock.coordinator.dt_util.utcnow",
+            return_value=now,
         ),
         patch(
             "homeassistant.components.roborock.coordinator.RoborockMqttClientV1.get_map_v1",
@@ -243,7 +259,10 @@ 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
+    # The map should load fine from the coordinator, but it should not update the
+    # last_updated timestamp.
+    assert resp.ok
+    assert previous_state == hass.states.get("image.roborock_s7_maxv_upstairs").state
 
 
 async def test_index_error_map(
@@ -257,6 +276,7 @@ async def test_index_error_map(
     # Copy the device prop so we don't override it
     prop = copy.deepcopy(PROP)
     prop.status.in_cleaning = 1
+    previous_state = hass.states.get("image.roborock_s7_maxv_upstairs").state
     # Update image, but get IndexError for image.
     with (
         patch(
@@ -268,9 +288,13 @@ async def test_index_error_map(
             return_value=prop,
         ),
         patch(
-            "homeassistant.components.roborock.image.dt_util.utcnow", return_value=now
+            "homeassistant.components.roborock.coordinator.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
+    # The map should load fine from the coordinator, but it should not update the
+    # last_updated timestamp.
+    assert resp.ok
+    assert previous_state == hass.states.get("image.roborock_s7_maxv_upstairs").state

tests/components/roborock/test_init.py

@@ -189,7 +189,7 @@ async def test_remove_from_hass(
     await hass.config_entries.async_unload(setup_entry.entry_id)
     assert cleanup_map_storage.exists()
     paths = list(cleanup_map_storage.walk())
-    assert len(paths) == 3  # One map image and two directories
+    assert len(paths) == 4  # Two map image and two directories
 
     await hass.config_entries.async_remove(setup_entry.entry_id)
     # After removal, directories should be empty.
@@ -219,7 +219,7 @@ async def test_oserror_remove_image(
 
     assert cleanup_map_storage.exists()
     paths = list(cleanup_map_storage.walk())
-    assert len(paths) == 3  # One map image and two directories
+    assert len(paths) == 4  # Two map image and two directories
 
     with patch(
         "homeassistant.components.roborock.roborock_storage.shutil.rmtree",
@@ -242,7 +242,7 @@ async def test_not_supported_protocol(
         "homeassistant.components.roborock.RoborockApiClient.get_home_data_v2",
         return_value=home_data_copy,
     ):
-        await async_setup_component(hass, DOMAIN, {})
+        await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
         await hass.async_block_till_done()
     assert "because its protocol version random" in caplog.text