From 6f3be3e50548db071a05dfb8d8a65cce14c7cafa Mon Sep 17 00:00:00 2001 From: Luke Lashley Date: Wed, 7 Feb 2024 03:13:51 -0500 Subject: [PATCH] Move Roborock map retrieval to coordinator and made map always diagnostic (#104680) Co-authored-by: Robert Resch --- homeassistant/components/roborock/__init__.py | 5 ++ .../components/roborock/coordinator.py | 9 +++ homeassistant/components/roborock/image.py | 77 ++++++++----------- tests/components/roborock/conftest.py | 3 + tests/components/roborock/test_init.py | 14 ++++ 5 files changed, 64 insertions(+), 44 deletions(-) diff --git a/homeassistant/components/roborock/__init__.py b/homeassistant/components/roborock/__init__.py index e0dfb2b271f..f8ceb121fe4 100644 --- a/homeassistant/components/roborock/__init__.py +++ b/homeassistant/components/roborock/__init__.py @@ -122,6 +122,11 @@ async def setup_device( # Verify we can communicate locally - if we can't, switch to cloud api await coordinator.verify_api() coordinator.api.is_available = True + try: + await coordinator.get_maps() + except RoborockException as err: + _LOGGER.warning("Failed to get map data") + _LOGGER.debug(err) try: await coordinator.async_config_entry_first_refresh() except ConfigEntryNotReady as ex: diff --git a/homeassistant/components/roborock/coordinator.py b/homeassistant/components/roborock/coordinator.py index cd08cf871d4..3864a90b16d 100644 --- a/homeassistant/components/roborock/coordinator.py +++ b/homeassistant/components/roborock/coordinator.py @@ -59,6 +59,8 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]): if mac := self.roborock_device_info.network_info.mac: self.device_info[ATTR_CONNECTIONS] = {(dr.CONNECTION_NETWORK_MAC, mac)} + # Maps from map flag to map name + self.maps: dict[int, str] = {} async def verify_api(self) -> None: """Verify that the api is reachable. If it is not, switch clients.""" @@ -107,3 +109,10 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]): self.current_map = ( self.roborock_device_info.props.status.map_status - 3 ) // 4 + + async def get_maps(self) -> None: + """Add a map to the coordinators mapping.""" + maps = await self.api.get_multi_maps_list() + if maps and maps.map_info: + for roborock_map in maps.map_info: + self.maps[roborock_map.mapFlag] = roborock_map.name diff --git a/homeassistant/components/roborock/image.py b/homeassistant/components/roborock/image.py index b2a14b57819..66957232679 100644 --- a/homeassistant/components/roborock/image.py +++ b/homeassistant/components/roborock/image.py @@ -66,13 +66,7 @@ class RoborockMap(RoborockCoordinatedEntity, ImageEntity): self._attr_image_last_updated = dt_util.utcnow() self.map_flag = map_flag self.cached_map = self._create_image(starting_map) - - @property - def entity_category(self) -> EntityCategory | None: - """Return diagnostic entity category for any non-selected maps.""" - if not self.is_selected: - return EntityCategory.DIAGNOSTIC - return None + self._attr_entity_category = EntityCategory.DIAGNOSTIC @property def is_selected(self) -> bool: @@ -127,42 +121,37 @@ async def create_coordinator_maps( Only one map can be loaded at a time per device. """ entities = [] - maps = await coord.cloud_api.get_multi_maps_list() - if maps is not None and maps.map_info is not None: - 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( - maps.map_info, key=lambda data: data.mapFlag == cur_map, reverse=True + + 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_name 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]) + # We cannot get the map until the roborock servers fully process the + # map change. + await asyncio.sleep(MAP_SLEEP) + # Get the map data + api_data: bytes = await coord.cloud_api.get_map_v1() + entities.append( + RoborockMap( + f"{slugify(coord.roborock_device_info.device.duid)}_map_{map_name}", + coord, + map_flag, + api_data, + map_name, + ) ) - for roborock_map in maps_info: - # Load the map - so we can access it with get_map_v1 - if roborock_map.mapFlag != cur_map: - # Only change the map and sleep if we have multiple maps. - await coord.api.send_command( - RoborockCommand.LOAD_MULTI_MAP, [roborock_map.mapFlag] - ) - # We cannot get the map until the roborock servers fully process the - # map change. - await asyncio.sleep(MAP_SLEEP) - # Get the map data - api_data: bytes = await coord.cloud_api.get_map_v1() - entities.append( - RoborockMap( - f"{slugify(coord.roborock_device_info.device.duid)}_map_{roborock_map.name}", - coord, - roborock_map.mapFlag, - api_data, - roborock_map.name, - ) - ) - if len(maps.map_info) != 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] - ) + 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]) return entities diff --git a/tests/components/roborock/conftest.py b/tests/components/roborock/conftest.py index 711ae203e0f..efbc2ea7f9d 100644 --- a/tests/components/roborock/conftest.py +++ b/tests/components/roborock/conftest.py @@ -45,6 +45,9 @@ def bypass_api_fixture() -> None: ), patch( "homeassistant.components.roborock.coordinator.RoborockMqttClient.get_multi_maps_list", return_value=MULTI_MAP_LIST, + ), patch( + "homeassistant.components.roborock.coordinator.RoborockLocalClient.get_multi_maps_list", + return_value=MULTI_MAP_LIST, ), patch( "homeassistant.components.roborock.image.RoborockMapDataParser.parse", return_value=MAP_DATA, diff --git a/tests/components/roborock/test_init.py b/tests/components/roborock/test_init.py index 5d1afaf8f84..608263a3496 100644 --- a/tests/components/roborock/test_init.py +++ b/tests/components/roborock/test_init.py @@ -107,6 +107,20 @@ async def test_local_client_fails_props( assert mock_roborock_entry.state is ConfigEntryState.SETUP_RETRY +async def test_fails_maps_continue( + hass: HomeAssistant, mock_roborock_entry: MockConfigEntry, bypass_api_fixture +) -> None: + """Test that if we fail to get the maps, we still setup.""" + with patch( + "homeassistant.components.roborock.coordinator.RoborockLocalClient.get_multi_maps_list", + side_effect=RoborockException(), + ): + await async_setup_component(hass, DOMAIN, {}) + assert mock_roborock_entry.state is ConfigEntryState.LOADED + # No map data means no images + assert len(hass.states.async_all("image")) == 0 + + async def test_reauth_started( hass: HomeAssistant, bypass_api_fixture, mock_roborock_entry: MockConfigEntry ) -> None: