Add ability to cache Roborock maps instead of always reloading (#112047)

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>
Co-authored-by: Allen Porter <allen.porter@gmail.com>
Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>
Co-authored-by: Allen Porter <allen@thebends.org>
Co-authored-by: Robert Resch <robert@resch.dev>

homeassistant/components/roborock/__init__.py

@@ -28,6 +28,7 @@ from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 
 from .const import CONF_BASE_URL, CONF_USER_DATA, DOMAIN, PLATFORMS
 from .coordinator import RoborockDataUpdateCoordinator, RoborockDataUpdateCoordinatorA01
+from .roborock_storage import async_remove_map_storage
 
 SCAN_INTERVAL = timedelta(seconds=30)
 
@@ -259,3 +260,8 @@ async def update_listener(hass: HomeAssistant, entry: RoborockConfigEntry) -> No
     """Handle options update."""
     # Reload entry to update data
     await hass.config_entries.async_reload(entry.entry_id)
+
+
+async def async_remove_entry(hass: HomeAssistant, entry: RoborockConfigEntry) -> None:
+    """Handle removal of an entry."""
+    await async_remove_map_storage(hass, entry.entry_id)

homeassistant/components/roborock/const.py

@@ -49,5 +49,7 @@ IMAGE_CACHE_INTERVAL = 90
 MAP_SLEEP = 3
 
 GET_MAPS_SERVICE_NAME = "get_maps"
+MAP_FILE_FORMAT = "PNG"
+MAP_FILENAME_SUFFIX = ".png"
 SET_VACUUM_GOTO_POSITION_SERVICE_NAME = "set_vacuum_goto_position"
 GET_VACUUM_CURRENT_POSITION_SERVICE_NAME = "get_vacuum_current_position"

homeassistant/components/roborock/coordinator.py

@@ -16,6 +16,7 @@ from roborock.version_1_apis.roborock_local_client_v1 import RoborockLocalClient
 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
 from roborock.version_a01_apis import RoborockClientA01
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_CONNECTIONS
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
@@ -26,6 +27,7 @@ from homeassistant.util import slugify
 
 from .const import DOMAIN
 from .models import RoborockA01HassDeviceInfo, RoborockHassDeviceInfo, RoborockMapInfo
+from .roborock_storage import RoborockMapStorage
 
 SCAN_INTERVAL = timedelta(seconds=30)
 
@@ -35,6 +37,8 @@ _LOGGER = logging.getLogger(__name__)
 class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
     """Class to manage fetching data from the API."""
 
+    config_entry: ConfigEntry
+
     def __init__(
         self,
         hass: HomeAssistant,
@@ -72,6 +76,9 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
         # Maps from map flag to map name
         self.maps: dict[int, RoborockMapInfo] = {}
         self._home_data_rooms = {str(room.id): room.name for room in home_data_rooms}
+        self.map_storage = RoborockMapStorage(
+            hass, self.config_entry.entry_id, slugify(self.duid)
+        )
 
     async def _async_setup(self) -> None:
         """Set up the coordinator."""

homeassistant/components/roborock/image.py

@@ -1,26 +1,33 @@
 """Support for Roborock image."""
 
 import asyncio
+from collections.abc import Callable
 from datetime import datetime
 import io
-from itertools import chain
 
 from roborock import RoborockCommand
 from vacuum_map_parser_base.config.color import ColorsPalette
-from vacuum_map_parser_base.config.drawable import Drawable
 from vacuum_map_parser_base.config.image_config import ImageConfig
 from vacuum_map_parser_base.config.size import Sizes
 from vacuum_map_parser_roborock.map_data_parser import RoborockMapDataParser
 
 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 AddEntitiesCallback
-from homeassistant.util import dt as dt_util, slugify
+from homeassistant.util import dt as dt_util
 
 from . import RoborockConfigEntry
-from .const import DEFAULT_DRAWABLES, DOMAIN, DRAWABLES, IMAGE_CACHE_INTERVAL, MAP_SLEEP
+from .const import (
+    DEFAULT_DRAWABLES,
+    DOMAIN,
+    DRAWABLES,
+    IMAGE_CACHE_INTERVAL,
+    MAP_FILE_FORMAT,
+    MAP_SLEEP,
+)
 from .coordinator import RoborockDataUpdateCoordinator
 from .entity import RoborockCoordinatedEntityV1
 
@@ -37,17 +44,35 @@ async def async_setup_entry(
         for drawable, default_value in DEFAULT_DRAWABLES.items()
         if config_entry.options.get(DRAWABLES, {}).get(drawable, default_value)
     ]
-    entities = list(
+    parser = RoborockMapDataParser(
-        chain.from_iterable(
+        ColorsPalette(), Sizes(), drawables, ImageConfig(), []
+    )
+
+    def parse_image(map_bytes: bytes) -> bytes | None:
+        parsed_map = parser.parse(map_bytes)
+        if parsed_map.image is None:
+            return None
+        img_byte_arr = io.BytesIO()
+        parsed_map.image.data.save(img_byte_arr, format=MAP_FILE_FORMAT)
+        return img_byte_arr.getvalue()
+
     await asyncio.gather(
-                *(
+        *(refresh_coordinators(hass, coord) for coord in config_entry.runtime_data.v1)
-                    create_coordinator_maps(coord, drawables)
+    )
+    async_add_entities(
+        (
+            RoborockMap(
+                config_entry,
+                f"{coord.duid_slug}_map_{map_info.name}",
+                coord,
+                map_info.flag,
+                map_info.name,
+                parse_image,
+            )
             for coord in config_entry.runtime_data.v1
+            for map_info in coord.maps.values()
+        ),
     )
-            )
-        )
-    )
-    async_add_entities(entities)
 
 
 class RoborockMap(RoborockCoordinatedEntityV1, ImageEntity):
@@ -55,39 +80,27 @@ class RoborockMap(RoborockCoordinatedEntityV1, ImageEntity):
 
     _attr_has_entity_name = True
     image_last_updated: datetime
+    _attr_name: str
 
     def __init__(
         self,
+        config_entry: ConfigEntry,
         unique_id: str,
         coordinator: RoborockDataUpdateCoordinator,
         map_flag: int,
-        starting_map: bytes,
         map_name: str,
-        drawables: list[Drawable],
+        parser: Callable[[bytes], bytes | None],
     ) -> None:
         """Initialize a Roborock map."""
         RoborockCoordinatedEntityV1.__init__(self, unique_id, coordinator)
         ImageEntity.__init__(self, coordinator.hass)
+        self.config_entry = config_entry
         self._attr_name = map_name
-        self.parser = RoborockMapDataParser(
+        self.parser = parser
-            ColorsPalette(), Sizes(), drawables, ImageConfig(), []
-        )
-        self._attr_image_last_updated = dt_util.utcnow()
         self.map_flag = map_flag
-        try:
-            self.cached_map = self._create_image(starting_map)
-        except HomeAssistantError:
-            # If we failed to update the image on init,
-            # we set cached_map to empty bytes
-            # so that we are unavailable and can try again later.
         self.cached_map = b""
         self._attr_entity_category = EntityCategory.DIAGNOSTIC
 
-    @property
-    def available(self) -> bool:
-        """Determines if the entity is available."""
-        return self.cached_map != b""
-
     @property
     def is_selected(self) -> bool:
         """Return if this map is the currently selected map."""
@@ -106,6 +119,14 @@ class RoborockMap(RoborockCoordinatedEntityV1, ImageEntity):
             and bool(self.coordinator.roborock_device_info.props.status.in_cleaning)
         )
 
+    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.cached_map = content or b""
+        self._attr_image_last_updated = dt_util.utcnow()
+        self.async_write_ha_state()
+
     def _handle_coordinator_update(self) -> None:
         # Bump last updated every third time the coordinator runs, so that async_image
         # will be called and we will evaluate on the new coordinator data if we should
@@ -126,47 +147,40 @@ class RoborockMap(RoborockCoordinatedEntityV1, ImageEntity):
                 ),
                 return_exceptions=True,
             )
-            if not isinstance(response[0], bytes):
+            if (
+                not isinstance(response[0], bytes)
+                or (content := self.parser(response[0])) is None
+            ):
                 raise HomeAssistantError(
                     translation_domain=DOMAIN,
                     translation_key="map_failure",
                 )
-            map_data = response[0]
+            if self.cached_map != content:
-            self.cached_map = self._create_image(map_data)
+                self.cached_map = content
+                self.config_entry.async_create_task(
+                    self.hass,
+                    self.coordinator.map_storage.async_save_map(
+                        self.map_flag,
+                        content,
+                    ),
+                    f"{self.unique_id} map",
+                )
         return self.cached_map
 
-    def _create_image(self, map_bytes: bytes) -> bytes:
-        """Create an image using the map parser."""
-        parsed_map = self.parser.parse(map_bytes)
-        if parsed_map.image is None:
-            raise HomeAssistantError(
-                translation_domain=DOMAIN,
-                translation_key="map_failure",
-            )
-        img_byte_arr = io.BytesIO()
-        parsed_map.image.data.save(img_byte_arr, format="PNG")
-        return img_byte_arr.getvalue()
 
-
+async def refresh_coordinators(
-async def create_coordinator_maps(
+    hass: HomeAssistant, coord: RoborockDataUpdateCoordinator
-    coord: RoborockDataUpdateCoordinator, drawables: list[Drawable]
+) -> None:
-) -> list[RoborockMap]:
     """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.
     """
-    entities = []
     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
+    map_flags = sorted(coord.maps, key=lambda data: data == cur_map, reverse=True)
-    # load_multi_map call.
+    for map_flag in map_flags:
-    maps_info = sorted(
-        coord.maps.items(), key=lambda data: data[0] == cur_map, reverse=True
-    )
-    for map_flag, map_info 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])
@@ -174,28 +188,11 @@ async def create_coordinator_maps(
             # We cannot get the map until the roborock servers fully process the
             # map change.
             await asyncio.sleep(MAP_SLEEP)
-        # Get the map data
+        await coord.set_current_map_rooms()
-        map_update = await asyncio.gather(
+
-            *[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.
-        api_data: bytes = map_update[0] if isinstance(map_update[0], bytes) else b""
-        entities.append(
-            RoborockMap(
-                f"{slugify(coord.duid)}_map_{map_info.name}",
-                coord,
-                map_flag,
-                api_data,
-                map_info.name,
-                drawables,
-            )
-        )
     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
-    return entities

homeassistant/components/roborock/roborock_storage.py

@@ -0,0 +1,81 @@
+"""Roborock storage."""
+
+import logging
+from pathlib import Path
+import shutil
+
+from homeassistant.core import HomeAssistant
+
+from .const import DOMAIN, MAP_FILENAME_SUFFIX
+
+_LOGGER = logging.getLogger(__name__)
+
+STORAGE_PATH = f".storage/{DOMAIN}"
+MAPS_PATH = "maps"
+
+
+def _storage_path_prefix(hass: HomeAssistant, entry_id: str) -> Path:
+    return Path(hass.config.path(STORAGE_PATH)) / entry_id
+
+
+class RoborockMapStorage:
+    """Store and retrieve maps for a Roborock device.
+
+    An instance of RoborockMapStorage is created for each device and manages
+    local storage of maps for that device.
+    """
+
+    def __init__(self, hass: HomeAssistant, entry_id: str, device_id_slug: str) -> None:
+        """Initialize RoborockMapStorage."""
+        self._hass = hass
+        self._path_prefix = (
+            _storage_path_prefix(hass, entry_id) / MAPS_PATH / device_id_slug
+        )
+
+    async def async_load_map(self, map_flag: int) -> bytes | None:
+        """Load maps from disk."""
+        filename = self._path_prefix / f"{map_flag}{MAP_FILENAME_SUFFIX}"
+        return await self._hass.async_add_executor_job(self._load_map, filename)
+
+    def _load_map(self, filename: Path) -> bytes | None:
+        """Load maps from disk."""
+        if not filename.exists():
+            return None
+        try:
+            return filename.read_bytes()
+        except OSError as err:
+            _LOGGER.debug("Unable to read map file: %s %s", filename, err)
+            return None
+
+    async def async_save_map(self, map_flag: int, content: bytes) -> None:
+        """Write map if it should be updated."""
+        filename = self._path_prefix / f"{map_flag}{MAP_FILENAME_SUFFIX}"
+        await self._hass.async_add_executor_job(self._save_map, filename, content)
+
+    def _save_map(self, filename: Path, content: bytes) -> None:
+        """Write the map to disk."""
+        _LOGGER.debug("Saving map to disk: %s", filename)
+        try:
+            filename.parent.mkdir(parents=True, exist_ok=True)
+        except OSError as err:
+            _LOGGER.error("Unable to create map directory: %s %s", filename, err)
+            return
+        try:
+            filename.write_bytes(content)
+        except OSError as err:
+            _LOGGER.error("Unable to write map file: %s %s", filename, err)
+
+
+async def async_remove_map_storage(hass: HomeAssistant, entry_id: str) -> None:
+    """Remove all map storage  associated with a config entry."""
+
+    def remove(path_prefix: Path) -> None:
+        try:
+            if path_prefix.exists():
+                shutil.rmtree(path_prefix, ignore_errors=True)
+        except OSError as err:
+            _LOGGER.error("Unable to remove map files in %s: %s", path_prefix, err)
+
+    path_prefix = _storage_path_prefix(hass, entry_id)
+    _LOGGER.debug("Removing maps from disk store: %s", path_prefix)
+    await hass.async_add_executor_job(remove, path_prefix)

tests/components/roborock/conftest.py

@@ -2,8 +2,11 @@
 
 from collections.abc import Generator
 from copy import deepcopy
+import pathlib
+import shutil
 from typing import Any
 from unittest.mock import Mock, patch
+import uuid
 
 import pytest
 from roborock import RoborockCategory, RoomMapping
@@ -70,6 +73,9 @@ def bypass_api_fixture() -> None:
     with (
         patch("homeassistant.components.roborock.RoborockMqttClientV1.async_connect"),
         patch("homeassistant.components.roborock.RoborockMqttClientV1._send_command"),
+        patch(
+            "homeassistant.components.roborock.coordinator.RoborockMqttClientV1._send_command"
+        ),
         patch(
             "homeassistant.components.roborock.RoborockApiClient.get_home_data_v2",
             return_value=HOME_DATA,
@@ -196,6 +202,7 @@ async def setup_entry(
     hass: HomeAssistant,
     bypass_api_fixture,
     mock_roborock_entry: MockConfigEntry,
+    cleanup_map_storage: pathlib.Path,
     platforms: list[Platform],
 ) -> Generator[MockConfigEntry]:
     """Set up the Roborock platform."""
@@ -203,3 +210,19 @@ async def setup_entry(
         assert await async_setup_component(hass, DOMAIN, {})
         await hass.async_block_till_done()
         yield mock_roborock_entry
+
+
+@pytest.fixture
+def cleanup_map_storage(
+    hass: HomeAssistant, mock_roborock_entry: MockConfigEntry
+) -> Generator[pathlib.Path]:
+    """Test cleanup, remove any map storage persisted during the test."""
+    tmp_path = str(uuid.uuid4())
+    with patch(
+        "homeassistant.components.roborock.roborock_storage.STORAGE_PATH", new=tmp_path
+    ):
+        storage_path = (
+            pathlib.Path(hass.config.path(tmp_path)) / mock_roborock_entry.entry_id
+        )
+        yield storage_path
+        shutil.rmtree(str(storage_path), ignore_errors=True)

tests/components/roborock/test_image.py

@@ -3,13 +3,16 @@
 import copy
 from datetime import timedelta
 from http import HTTPStatus
+import io
 from unittest.mock import patch
 
+from PIL import Image
 import pytest
 from roborock import RoborockException
+from vacuum_map_parser_base.map_data import ImageConfig, ImageData
 
 from homeassistant.components.roborock import DOMAIN
-from homeassistant.const import STATE_UNAVAILABLE, Platform
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 from homeassistant.util import dt as dt_util
@@ -32,22 +35,27 @@ async def test_floorplan_image(
     hass_client: ClientSessionGenerator,
 ) -> None:
     """Test floor plan map image is correctly set up."""
-    # Setup calls the image parsing the first time and caches it.
     assert len(hass.states.async_all("image")) == 4
 
     assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None
-    # call a second time -should return cached data
+    # Load the image on demand
     client = await hass_client()
     resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
     assert resp.status == HTTPStatus.OK
     body = await resp.read()
     assert body is not None
-    # Call a third time - this time forcing it to update
+    assert body[0:4] == b"\x89PNG"
-    now = dt_util.utcnow() + timedelta(seconds=91)
+
+    # Call a second time - this time forcing it to update - and save new image
+    now = dt_util.utcnow() + timedelta(minutes=61)
 
     # Copy the device prop so we don't override it
     prop = copy.deepcopy(PROP)
     prop.status.in_cleaning = 1
+    new_map_data = copy.deepcopy(MAP_DATA)
+    new_map_data.image = ImageData(
+        100, 10, 10, 10, 10, ImageConfig(), Image.new("RGB", (2, 2)), lambda p: p
+    )
     with (
         patch(
             "homeassistant.components.roborock.coordinator.RoborockLocalClientV1.get_prop",
@@ -56,6 +64,10 @@ async def test_floorplan_image(
         patch(
             "homeassistant.components.roborock.image.dt_util.utcnow", return_value=now
         ),
+        patch(
+            "homeassistant.components.roborock.image.RoborockMapDataParser.parse",
+            return_value=new_map_data,
+        ) as parse_map,
     ):
         async_fire_time_changed(hass, now)
         await hass.async_block_till_done()
@@ -63,6 +75,7 @@ async def test_floorplan_image(
     assert resp.status == HTTPStatus.OK
     body = await resp.read()
     assert body is not None
+    assert parse_map.call_count == 1
 
 
 async def test_floorplan_image_failed_parse(
@@ -97,13 +110,101 @@ async def test_floorplan_image_failed_parse(
     assert not resp.ok
 
 
+async def test_load_stored_image(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    setup_entry: MockConfigEntry,
+) -> None:
+    """Test that we correctly load an image from storage when it already exists."""
+    img_byte_arr = io.BytesIO()
+    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
+    client = await hass_client()
+    resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
+    assert resp.status == HTTPStatus.OK
+
+    with patch(
+        "homeassistant.components.roborock.image.RoborockMapDataParser.parse",
+    ) as parse_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)
+        await hass.async_block_till_done()
+        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
+        body = await resp.read()
+        assert body == img_bytes
+
+        # Ensure that we never tried to update the map, and only used the cached image.
+        assert parse_map.call_count == 0
+
+
+async def test_fail_to_save_image(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    mock_roborock_entry: MockConfigEntry,
+    bypass_api_fixture,
+    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.
+    with patch(
+        "homeassistant.components.roborock.roborock_storage.Path.write_bytes",
+        side_effect=OSError,
+    ):
+        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
+
+    assert "Unable to write map file" in caplog.text
+
+
+async def test_fail_to_load_image(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    setup_entry: MockConfigEntry,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that we gracefully handle failing to load an image."""
+    with (
+        patch(
+            "homeassistant.components.roborock.image.RoborockMapDataParser.parse",
+        ) as parse_map,
+        patch(
+            "homeassistant.components.roborock.roborock_storage.Path.exists",
+            return_value=True,
+        ),
+        patch(
+            "homeassistant.components.roborock.roborock_storage.Path.read_bytes",
+            side_effect=OSError,
+        ) as read_bytes,
+    ):
+        # 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)
+        await hass.async_block_till_done()
+        assert read_bytes.call_count == 4
+        # Ensure that we never updated the map manually since we couldn't load it.
+        assert parse_map.call_count == 0
+    assert "Unable to read map file" in caplog.text
+
+
 async def test_fail_parse_on_startup(
     hass: HomeAssistant,
     hass_client: ClientSessionGenerator,
     mock_roborock_entry: MockConfigEntry,
     bypass_api_fixture,
 ) -> None:
-    """Test that if we fail parsing on startup, we create the entity but set it as unavailable."""
+    """Test that if we fail parsing on startup, we still create the entity."""
     map_data = copy.deepcopy(MAP_DATA)
     map_data.image = None
     with patch(
@@ -115,7 +216,28 @@ async def test_fail_parse_on_startup(
     assert (
         image_entity := hass.states.get("image.roborock_s7_maxv_upstairs")
     ) is not None
-    assert image_entity.state == STATE_UNAVAILABLE
+    assert image_entity.state
+
+
+async def test_fail_get_map_on_startup(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    mock_roborock_entry: MockConfigEntry,
+    bypass_api_fixture,
+) -> None:
+    """Test that if we fail getting map on startup, we can still create the entity."""
+    with (
+        patch(
+            "homeassistant.components.roborock.coordinator.RoborockMqttClientV1.get_map_v1",
+            return_value=None,
+        ),
+    ):
+        await async_setup_component(hass, DOMAIN, {})
+        await hass.async_block_till_done()
+    assert (
+        image_entity := hass.states.get("image.roborock_s7_maxv_upstairs")
+    ) is not None
+    assert image_entity.state
 
 
 async def test_fail_updating_image(

tests/components/roborock/test_init.py

@@ -1,6 +1,8 @@
 """Test for Roborock init."""
 
 from copy import deepcopy
+from http import HTTPStatus
+import pathlib
 from unittest.mock import patch
 
 import pytest
@@ -13,12 +15,14 @@ from roborock import (
 
 from homeassistant.components.roborock.const 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
 
 from .mock_data import HOME_DATA
 
 from tests.common import MockConfigEntry
+from tests.typing import ClientSessionGenerator
 
 
 async def test_unload_entry(
@@ -163,6 +167,60 @@ async def test_reauth_started(
     assert flows[0]["step_id"] == "reauth_confirm"
 
 
+@pytest.mark.parametrize("platforms", [[Platform.IMAGE]])
+async def test_remove_from_hass(
+    hass: HomeAssistant,
+    bypass_api_fixture,
+    setup_entry: MockConfigEntry,
+    hass_client: ClientSessionGenerator,
+    cleanup_map_storage: pathlib.Path,
+) -> None:
+    """Test that removing from hass removes any existing images."""
+
+    # Ensure some image content is cached
+    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")
+    assert resp.status == HTTPStatus.OK
+
+    assert cleanup_map_storage.exists()
+    paths = list(cleanup_map_storage.walk())
+    assert len(paths) == 3  # One map image and two directories
+
+    await hass.config_entries.async_remove(setup_entry.entry_id)
+    # After removal, directories should be empty.
+    assert not cleanup_map_storage.exists()
+
+
+@pytest.mark.parametrize("platforms", [[Platform.IMAGE]])
+async def test_oserror_remove_image(
+    hass: HomeAssistant,
+    bypass_api_fixture,
+    setup_entry: MockConfigEntry,
+    cleanup_map_storage: pathlib.Path,
+    hass_client: ClientSessionGenerator,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that we gracefully handle failing to remove an image."""
+
+    # Ensure some image content is cached
+    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")
+    assert resp.status == HTTPStatus.OK
+
+    assert cleanup_map_storage.exists()
+    paths = list(cleanup_map_storage.walk())
+    assert len(paths) == 3  # One map image and two directories
+
+    with patch(
+        "homeassistant.components.roborock.roborock_storage.shutil.rmtree",
+        side_effect=OSError,
+    ):
+        await hass.config_entries.async_remove(setup_entry.entry_id)
+    assert "Unable to remove map files" in caplog.text
+
+
 async def test_not_supported_protocol(
     hass: HomeAssistant,
     bypass_api_fixture,