Persist roborock maps to disk only on shutdown (#136889)

* Persist roborock maps to disk only on shutdown

* Rename on_unload to on_stop

* Spawn 1 executor thread and block writes to disk

* Update tests/components/roborock/test_image.py

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

* Use config entry setup instead of component setup

---------

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

homeassistant/components/roborock/__init__.py

@@ -22,7 +22,7 @@ from roborock.version_a01_apis import RoborockMqttClientA01
 from roborock.web_api import RoborockApiClient
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_USERNAME
+from homeassistant.const import CONF_USERNAME, EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 
@@ -118,13 +118,21 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) ->
         )
     valid_coordinators = RoborockCoordinators(v1_coords, a01_coords)
 
-    async def on_unload() -> None:
+    async def on_stop(_: Any) -> None:
-        release_tasks = set()
+        _LOGGER.debug("Shutting down roborock")
-        for coordinator in valid_coordinators.values():
+        await asyncio.gather(
-            release_tasks.add(coordinator.release())
+            *(
-        await asyncio.gather(*release_tasks)
+                coordinator.async_shutdown()
+                for coordinator in valid_coordinators.values()
+            )
+        )
 
-    entry.async_on_unload(on_unload)
+    entry.async_on_unload(
+        hass.bus.async_listen_once(
+            EVENT_HOMEASSISTANT_STOP,
+            on_stop,
+        )
+    )
     entry.runtime_data = valid_coordinators
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
@@ -209,7 +217,7 @@ async def setup_device_v1(
     try:
         await coordinator.async_config_entry_first_refresh()
     except ConfigEntryNotReady as ex:
-        await coordinator.release()
+        await coordinator.async_shutdown()
         if isinstance(coordinator.api, RoborockMqttClientV1):
             _LOGGER.warning(
                 "Not setting up %s because the we failed to get data for the first time using the online client. "

homeassistant/components/roborock/coordinator.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from datetime import timedelta
 import logging
 
@@ -116,10 +117,14 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
                 # Right now this should never be called if the cloud api is the primary api,
                 # but in the future if it is, a new else should be added.
 
-    async def release(self) -> None:
+    async def async_shutdown(self) -> None:
-        """Disconnect from API."""
+        """Shutdown the coordinator."""
-        await self.api.async_release()
+        await super().async_shutdown()
-        await self.cloud_api.async_release()
+        await asyncio.gather(
+            self.map_storage.flush(),
+            self.api.async_release(),
+            self.cloud_api.async_release(),
+        )
 
     async def _update_device_prop(self) -> None:
         """Update device properties."""
@@ -226,8 +231,9 @@ class RoborockDataUpdateCoordinatorA01(
     ) -> dict[RoborockDyadDataProtocol | RoborockZeoProtocol, StateType]:
         return await self.api.update_values(self.request_protocols)
 
-    async def release(self) -> None:
+    async def async_shutdown(self) -> None:
-        """Disconnect from API."""
+        """Shutdown the coordinator on config entry unload."""
+        await super().async_shutdown()
         await self.api.async_release()
 
     @cached_property

homeassistant/components/roborock/image.py

@@ -157,13 +157,9 @@ class RoborockMap(RoborockCoordinatedEntityV1, ImageEntity):
                 )
             if self.cached_map != content:
                 self.cached_map = content
-                self.config_entry.async_create_task(
+                await self.coordinator.map_storage.async_save_map(
-                    self.hass,
+                    self.map_flag,
-                    self.coordinator.map_storage.async_save_map(
+                    content,
-                        self.map_flag,
-                        content,
-                    ),
-                    f"{self.unique_id} map",
                 )
         return self.cached_map
 

homeassistant/components/roborock/roborock_storage.py

@@ -31,6 +31,7 @@ class RoborockMapStorage:
         self._path_prefix = (
             _storage_path_prefix(hass, entry_id) / MAPS_PATH / device_id_slug
         )
+        self._write_queue: dict[int, bytes] = {}
 
     async def async_load_map(self, map_flag: int) -> bytes | None:
         """Load maps from disk."""
@@ -48,9 +49,22 @@ class RoborockMapStorage:
             return None
 
     async def async_save_map(self, map_flag: int, content: bytes) -> None:
-        """Write map if it should be updated."""
+        """Save the map to a pending write queue."""
-        filename = self._path_prefix / f"{map_flag}{MAP_FILENAME_SUFFIX}"
+        self._write_queue[map_flag] = content
-        await self._hass.async_add_executor_job(self._save_map, filename, content)
+
+    async def flush(self) -> None:
+        """Flush all maps to disk."""
+        _LOGGER.debug("Flushing %s maps to disk", len(self._write_queue))
+
+        queue = self._write_queue.copy()
+
+        def _flush_all() -> None:
+            for map_flag, content in queue.items():
+                filename = self._path_prefix / f"{map_flag}{MAP_FILENAME_SUFFIX}"
+                self._save_map(filename, content)
+
+        await self._hass.async_add_executor_job(_flush_all)
+        self._write_queue.clear()
 
     def _save_map(self, filename: Path, content: bytes) -> None:
         """Write the map to disk."""

tests/components/roborock/conftest.py

@@ -19,9 +19,9 @@ from homeassistant.components.roborock.const import (
     CONF_USER_DATA,
     DOMAIN,
 )
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.setup import async_setup_component
 
 from .mock_data import (
     BASE_URL,
@@ -207,13 +207,13 @@ async def setup_entry(
 ) -> Generator[MockConfigEntry]:
     """Set up the Roborock platform."""
     with patch("homeassistant.components.roborock.PLATFORMS", platforms):
-        assert await async_setup_component(hass, DOMAIN, {})
+        await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
         await hass.async_block_till_done()
         yield mock_roborock_entry
 
 
 @pytest.fixture
-def cleanup_map_storage(
+async def cleanup_map_storage(
     hass: HomeAssistant, mock_roborock_entry: MockConfigEntry
 ) -> Generator[pathlib.Path]:
     """Test cleanup, remove any map storage persisted during the test."""
@@ -225,4 +225,8 @@ def cleanup_map_storage(
             pathlib.Path(hass.config.path(tmp_path)) / mock_roborock_entry.entry_id
         )
         yield storage_path
+        # We need to first unload the config entry because unloading it will
+        # persist any unsaved maps to storage.
+        if mock_roborock_entry.state is ConfigEntryState.LOADED:
+            await hass.config_entries.async_unload(mock_roborock_entry.entry_id)
         shutil.rmtree(str(storage_path), ignore_errors=True)

tests/components/roborock/test_image.py

@@ -12,6 +12,7 @@ from roborock import RoborockException
 from vacuum_map_parser_base.map_data import ImageConfig, ImageData
 
 from homeassistant.components.roborock import DOMAIN
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
@@ -120,7 +121,7 @@ async def test_load_stored_image(
     MAP_DATA.image.data.save(img_byte_arr, format="PNG")
     img_bytes = img_byte_arr.getvalue()
 
-    # Load the image on demand, which should ensure it is cached on disk
+    # Load the image on demand, which should queue it to be cached on disk
     client = await hass_client()
     resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
     assert resp.status == HTTPStatus.OK
@@ -151,22 +152,25 @@ async def test_fail_to_save_image(
     caplog: pytest.LogCaptureFixture,
 ) -> None:
     """Test that we gracefully handle a oserror on saving an image."""
-    # Reload the config entry so that the map is saved in storage and entities exist.
+    await async_setup_component(hass, DOMAIN, {})
+    await hass.async_block_till_done()
+
+    # Ensure that map is still working properly.
+    assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None
+    client = await hass_client()
+    resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
+    # Test that we can get the image and it correctly serialized and unserialized.
+    assert resp.status == HTTPStatus.OK
+
     with patch(
         "homeassistant.components.roborock.roborock_storage.Path.write_bytes",
         side_effect=OSError,
     ):
-        await async_setup_component(hass, DOMAIN, {})
+        await hass.config_entries.async_unload(mock_roborock_entry.entry_id)
-        await hass.async_block_till_done()
+        assert "Unable to write map file" in caplog.text
 
-        # Ensure that map is still working properly.
+        # Config entry is unloaded successfully
-        assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None
+        assert mock_roborock_entry.state is ConfigEntryState.NOT_LOADED
-        client = await hass_client()
-        resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
-        # Test that we can get the image and it correctly serialized and unserialized.
-        assert resp.status == HTTPStatus.OK
-
-    assert "Unable to write map file" in caplog.text
 
 
 async def test_fail_to_load_image(

tests/components/roborock/test_init.py

@@ -183,6 +183,10 @@ async def test_remove_from_hass(
     resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
     assert resp.status == HTTPStatus.OK
 
+    assert not cleanup_map_storage.exists()
+
+    # Flush to disk
+    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
@@ -209,6 +213,10 @@ async def test_oserror_remove_image(
     resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
     assert resp.status == HTTPStatus.OK
 
+    # Image content is saved when unloading
+    assert not cleanup_map_storage.exists()
+    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