From a40bb2790ebb5f652b5a7112cb455e50d55ecd65 Mon Sep 17 00:00:00 2001 From: Luke Lashley Date: Sun, 16 Mar 2025 17:15:04 -0400 Subject: [PATCH] Move Roborock map refresh to coordinator (#140758) Move refresh coordinator to coordinator --- homeassistant/components/roborock/__init__.py | 3 ++ .../components/roborock/coordinator.py | 31 ++++++++++++++++ homeassistant/components/roborock/image.py | 37 +------------------ tests/components/roborock/conftest.py | 5 ++- tests/components/roborock/test_image.py | 30 +++++++++++++++ 5 files changed, 69 insertions(+), 37 deletions(-) diff --git a/homeassistant/components/roborock/__init__.py b/homeassistant/components/roborock/__init__.py index 955e50cd15b..1b90adaf6ec 100644 --- a/homeassistant/components/roborock/__init__.py +++ b/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") diff --git a/homeassistant/components/roborock/coordinator.py b/homeassistant/components/roborock/coordinator.py index 2f156545929..cbfd5e95a90 100644 --- a/homeassistant/components/roborock/coordinator.py +++ b/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[ diff --git a/homeassistant/components/roborock/image.py b/homeassistant/components/roborock/image.py index b56abaeebdb..382edbca744 100644 --- a/homeassistant/components/roborock/image.py +++ b/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 diff --git a/tests/components/roborock/conftest.py b/tests/components/roborock/conftest.py index cafac280620..b4fde5cc513 100644 --- a/tests/components/roborock/conftest.py +++ b/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( diff --git a/tests/components/roborock/test_image.py b/tests/components/roborock/test_image.py index d81f1289fe3..08f8ac504bf 100644 --- a/tests/components/roborock/test_image.py +++ b/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