From 4ce891512e14e6907e3f15a340b42bcde7a2d5ad Mon Sep 17 00:00:00 2001 From: Luke Lashley Date: Wed, 29 Jan 2025 12:16:28 -0500 Subject: [PATCH] Add ability to cache Roborock maps instead of always reloading (#112047) Co-authored-by: Paulus Schoutsen Co-authored-by: Allen Porter Co-authored-by: Joost Lekkerkerker Co-authored-by: Allen Porter Co-authored-by: Robert Resch --- homeassistant/components/roborock/__init__.py | 6 + homeassistant/components/roborock/const.py | 2 + .../components/roborock/coordinator.py | 7 + homeassistant/components/roborock/image.py | 151 +++++++++--------- .../components/roborock/roborock_storage.py | 81 ++++++++++ tests/components/roborock/conftest.py | 23 +++ tests/components/roborock/test_image.py | 136 +++++++++++++++- tests/components/roborock/test_init.py | 58 +++++++ 8 files changed, 380 insertions(+), 84 deletions(-) create mode 100644 homeassistant/components/roborock/roborock_storage.py diff --git a/homeassistant/components/roborock/__init__.py b/homeassistant/components/roborock/__init__.py index 9ab9226c9a5..1b34dc891d1 100644 --- a/homeassistant/components/roborock/__init__.py +++ b/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) diff --git a/homeassistant/components/roborock/const.py b/homeassistant/components/roborock/const.py index 4a9bd14bfe1..cc8d34fbadc 100644 --- a/homeassistant/components/roborock/const.py +++ b/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" diff --git a/homeassistant/components/roborock/coordinator.py b/homeassistant/components/roborock/coordinator.py index d34ba49da52..36333f1c55e 100644 --- a/homeassistant/components/roborock/coordinator.py +++ b/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.""" diff --git a/homeassistant/components/roborock/image.py b/homeassistant/components/roborock/image.py index 3818a039fb8..b0de4f9caa5 100644 --- a/homeassistant/components/roborock/image.py +++ b/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( - chain.from_iterable( - await asyncio.gather( - *( - create_coordinator_maps(coord, drawables) - for coord in config_entry.runtime_data.v1 - ) - ) - ) + parser = RoborockMapDataParser( + 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) + ) + 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( - ColorsPalette(), Sizes(), drawables, ImageConfig(), [] - ) - self._attr_image_last_updated = dt_util.utcnow() + self.parser = parser 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.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] - self.cached_map = self._create_image(map_data) + if self.cached_map != content: + 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 create_coordinator_maps( - coord: RoborockDataUpdateCoordinator, drawables: list[Drawable] -) -> list[RoborockMap]: +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. """ - 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 - # load_multi_map call. - 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 + 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]) @@ -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 - 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, - ) - ) + 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 - return entities diff --git a/homeassistant/components/roborock/roborock_storage.py b/homeassistant/components/roborock/roborock_storage.py new file mode 100644 index 00000000000..62e15e889be --- /dev/null +++ b/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) diff --git a/tests/components/roborock/conftest.py b/tests/components/roborock/conftest.py index 4df5f479b7c..e5fc5cb7eb6 100644 --- a/tests/components/roborock/conftest.py +++ b/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) diff --git a/tests/components/roborock/test_image.py b/tests/components/roborock/test_image.py index e240dccf7eb..90886f25929 100644 --- a/tests/components/roborock/test_image.py +++ b/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 - now = dt_util.utcnow() + timedelta(seconds=91) + assert body[0:4] == b"\x89PNG" + + # 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( diff --git a/tests/components/roborock/test_init.py b/tests/components/roborock/test_init.py index f4f490e68d9..efd1c3f66f4 100644 --- a/tests/components/roborock/test_init.py +++ b/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,