From 5ef42078a3462b718968842bbbb73a551c171b08 Mon Sep 17 00:00:00 2001 From: Luke Lashley Date: Mon, 8 Apr 2024 15:40:59 -0400 Subject: [PATCH] Add a service to get maps for Roborock (#111478) * add service to get rooms * fix linting * add snapshot * change map id --- homeassistant/components/roborock/__init__.py | 15 ++++++++--- homeassistant/components/roborock/const.py | 2 ++ .../components/roborock/coordinator.py | 27 ++++++++++++++++--- homeassistant/components/roborock/icons.json | 3 +++ homeassistant/components/roborock/image.py | 13 ++++++--- homeassistant/components/roborock/models.py | 9 +++++++ .../components/roborock/services.yaml | 4 +++ .../components/roborock/strings.json | 6 +++++ homeassistant/components/roborock/vacuum.py | 19 +++++++++++-- tests/components/roborock/conftest.py | 17 ++++++++++++ .../roborock/snapshots/test_vacuum.ambr | 27 +++++++++++++++++++ tests/components/roborock/test_vacuum.py | 19 +++++++++++++ 12 files changed, 147 insertions(+), 14 deletions(-) create mode 100644 homeassistant/components/roborock/services.yaml create mode 100644 tests/components/roborock/snapshots/test_vacuum.ambr diff --git a/homeassistant/components/roborock/__init__.py b/homeassistant/components/roborock/__init__.py index e64b83be1dd..141770e733d 100644 --- a/homeassistant/components/roborock/__init__.py +++ b/homeassistant/components/roborock/__init__.py @@ -8,7 +8,7 @@ from datetime import timedelta import logging from typing import Any -from roborock import RoborockException, RoborockInvalidCredentials +from roborock import HomeDataRoom, RoborockException, RoborockInvalidCredentials from roborock.cloud_api import RoborockMqttClient from roborock.containers import DeviceData, HomeDataDevice, HomeDataProduct, UserData from roborock.web_api import RoborockApiClient @@ -28,6 +28,7 @@ _LOGGER = logging.getLogger(__name__) async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up roborock from a config entry.""" + _LOGGER.debug("Integration async setup entry: %s", entry.as_dict()) user_data = UserData.from_dict(entry.data[CONF_USER_DATA]) @@ -56,7 +57,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: } # Get a Coordinator if the device is available or if we have connected to the device before coordinators = await asyncio.gather( - *build_setup_functions(hass, device_map, user_data, product_info), + *build_setup_functions( + hass, device_map, user_data, product_info, home_data.rooms + ), return_exceptions=True, ) # Valid coordinators are those where we had networking cached or we could get networking @@ -85,10 +88,13 @@ def build_setup_functions( device_map: dict[str, HomeDataDevice], user_data: UserData, product_info: dict[str, HomeDataProduct], + home_data_rooms: list[HomeDataRoom], ) -> list[Coroutine[Any, Any, RoborockDataUpdateCoordinator | None]]: """Create a list of setup functions that can later be called asynchronously.""" return [ - setup_device(hass, user_data, device, product_info[device.product_id]) + setup_device( + hass, user_data, device, product_info[device.product_id], home_data_rooms + ) for device in device_map.values() ] @@ -98,6 +104,7 @@ async def setup_device( user_data: UserData, device: HomeDataDevice, product_info: HomeDataProduct, + home_data_rooms: list[HomeDataRoom], ) -> RoborockDataUpdateCoordinator | None: """Set up a device Coordinator.""" mqtt_client = RoborockMqttClient(user_data, DeviceData(device, product_info.name)) @@ -117,7 +124,7 @@ async def setup_device( await mqtt_client.async_release() raise coordinator = RoborockDataUpdateCoordinator( - hass, device, networking, product_info, mqtt_client + hass, device, networking, product_info, mqtt_client, home_data_rooms ) # Verify we can communicate locally - if we can't, switch to cloud api await coordinator.verify_api() diff --git a/homeassistant/components/roborock/const.py b/homeassistant/components/roborock/const.py index 77f0be3363e..6b1ed975fca 100644 --- a/homeassistant/components/roborock/const.py +++ b/homeassistant/components/roborock/const.py @@ -30,3 +30,5 @@ IMAGE_DRAWABLES: list[Drawable] = [ IMAGE_CACHE_INTERVAL = 90 MAP_SLEEP = 3 + +GET_MAPS_SERVICE_NAME = "get_maps" diff --git a/homeassistant/components/roborock/coordinator.py b/homeassistant/components/roborock/coordinator.py index e682b119069..c5fd0c09c46 100644 --- a/homeassistant/components/roborock/coordinator.py +++ b/homeassistant/components/roborock/coordinator.py @@ -2,9 +2,11 @@ from __future__ import annotations +import asyncio from datetime import timedelta import logging +from roborock import HomeDataRoom from roborock.cloud_api import RoborockMqttClient from roborock.containers import DeviceData, HomeDataDevice, HomeDataProduct, NetworkInfo from roborock.exceptions import RoborockException @@ -18,7 +20,7 @@ from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from .const import DOMAIN -from .models import RoborockHassDeviceInfo +from .models import RoborockHassDeviceInfo, RoborockMapInfo SCAN_INTERVAL = timedelta(seconds=30) @@ -35,6 +37,7 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]): device_networking: NetworkInfo, product_info: HomeDataProduct, cloud_api: RoborockMqttClient, + home_data_rooms: list[HomeDataRoom], ) -> None: """Initialize.""" super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=SCAN_INTERVAL) @@ -61,7 +64,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] = {} + self.maps: dict[int, RoborockMapInfo] = {} + self._home_data_rooms = {str(room.id): room.name for room in home_data_rooms} async def verify_api(self) -> None: """Verify that the api is reachable. If it is not, switch clients.""" @@ -95,7 +99,7 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]): async def _async_update_data(self) -> DeviceProp: """Update data via library.""" try: - await self._update_device_prop() + await asyncio.gather(*(self._update_device_prop(), self.get_rooms())) self._set_current_map() except RoborockException as ex: raise UpdateFailed(ex) from ex @@ -117,4 +121,19 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]): 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 + self.maps[roborock_map.mapFlag] = RoborockMapInfo( + flag=roborock_map.mapFlag, name=roborock_map.name, rooms={} + ) + + async def get_rooms(self) -> None: + """Get all of the rooms for the current map.""" + # The api is only able to access rooms for the currently selected map + # So it is important this is only called when you have the map you care + # about selected. + if self.current_map in self.maps: + iot_rooms = await self.api.get_room_mapping() + if iot_rooms is not None: + for room in iot_rooms: + self.maps[self.current_map].rooms[room.segment_id] = ( + self._home_data_rooms.get(room.iot_id, "Unknown") + ) diff --git a/homeassistant/components/roborock/icons.json b/homeassistant/components/roborock/icons.json index 43e7f185433..babde739775 100644 --- a/homeassistant/components/roborock/icons.json +++ b/homeassistant/components/roborock/icons.json @@ -105,5 +105,8 @@ "default": "mdi:power-plug-off" } } + }, + "services": { + "get_maps": "mdi:floor-plan" } } diff --git a/homeassistant/components/roborock/image.py b/homeassistant/components/roborock/image.py index 3367f1b3017..775ab98fd59 100644 --- a/homeassistant/components/roborock/image.py +++ b/homeassistant/components/roborock/image.py @@ -130,23 +130,27 @@ async def create_coordinator_maps( maps_info = sorted( coord.maps.items(), key=lambda data: data[0] == cur_map, reverse=True ) - for map_flag, map_name in maps_info: + 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]) + coord.current_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() + map_update = await asyncio.gather( + *[coord.cloud_api.get_map_v1(), coord.get_rooms()] + ) + api_data: bytes = map_update[0] entities.append( RoborockMap( - f"{slugify(coord.roborock_device_info.device.duid)}_map_{map_name}", + f"{slugify(coord.roborock_device_info.device.duid)}_map_{map_info.name}", coord, map_flag, api_data, - map_name, + map_info.name, ) ) if len(coord.maps) != 1: @@ -154,4 +158,5 @@ async def create_coordinator_maps( # 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/models.py b/homeassistant/components/roborock/models.py index 45b98fddbc5..b516c0ee05c 100644 --- a/homeassistant/components/roborock/models.py +++ b/homeassistant/components/roborock/models.py @@ -24,3 +24,12 @@ class RoborockHassDeviceInfo: "product": self.product.as_dict(), "props": self.props.as_dict(), } + + +@dataclass +class RoborockMapInfo: + """A model to describe all information about a map we may want.""" + + flag: int + name: str + rooms: dict[int, str] diff --git a/homeassistant/components/roborock/services.yaml b/homeassistant/components/roborock/services.yaml new file mode 100644 index 00000000000..18de5c98c7b --- /dev/null +++ b/homeassistant/components/roborock/services.yaml @@ -0,0 +1,4 @@ +get_maps: + target: + entity: + domain: vacuum diff --git a/homeassistant/components/roborock/strings.json b/homeassistant/components/roborock/strings.json index 7c457a1935b..30aa64f626a 100644 --- a/homeassistant/components/roborock/strings.json +++ b/homeassistant/components/roborock/strings.json @@ -293,5 +293,11 @@ "no_coordinators": { "message": "No devices were able to successfully setup" } + }, + "services": { + "get_maps": { + "name": "Get maps", + "description": "Get the map and room information of your device." + } } } diff --git a/homeassistant/components/roborock/vacuum.py b/homeassistant/components/roborock/vacuum.py index 22d9353e2a2..d8108abf78c 100644 --- a/homeassistant/components/roborock/vacuum.py +++ b/homeassistant/components/roborock/vacuum.py @@ -1,5 +1,6 @@ """Support for Roborock vacuum class.""" +from dataclasses import asdict from typing import Any from roborock.code_mappings import RoborockStateCode @@ -17,11 +18,12 @@ from homeassistant.components.vacuum import ( VacuumEntityFeature, ) from homeassistant.config_entries import ConfigEntry -from homeassistant.core import HomeAssistant +from homeassistant.core import HomeAssistant, ServiceResponse, SupportsResponse +from homeassistant.helpers import entity_platform from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.util import slugify -from .const import DOMAIN +from .const import DOMAIN, GET_MAPS_SERVICE_NAME from .coordinator import RoborockDataUpdateCoordinator from .device import RoborockCoordinatedEntity @@ -66,6 +68,15 @@ async def async_setup_entry( for device_id, coordinator in coordinators.items() ) + platform = entity_platform.async_get_current_platform() + + platform.async_register_entity_service( + GET_MAPS_SERVICE_NAME, + {}, + RoborockVacuum.get_maps.__name__, + supports_response=SupportsResponse.ONLY, + ) + class RoborockVacuum(RoborockCoordinatedEntity, StateVacuumEntity): """General Representation of a Roborock vacuum.""" @@ -164,3 +175,7 @@ class RoborockVacuum(RoborockCoordinatedEntity, StateVacuumEntity): ) -> None: """Send a command to a vacuum cleaner.""" await self.send(command, params) + + async def get_maps(self) -> ServiceResponse: + """Get map information such as map id and room ids.""" + return {"maps": [asdict(map) for map in self.coordinator.maps.values()]} diff --git a/tests/components/roborock/conftest.py b/tests/components/roborock/conftest.py index 91331a1486a..2910fa38995 100644 --- a/tests/components/roborock/conftest.py +++ b/tests/components/roborock/conftest.py @@ -3,6 +3,7 @@ from unittest.mock import patch import pytest +from roborock import RoomMapping from homeassistant.components.roborock.const import ( CONF_BASE_URL, @@ -74,6 +75,22 @@ def bypass_api_fixture() -> None: "homeassistant.components.roborock.image.MAP_SLEEP", 0, ), + patch( + "homeassistant.components.roborock.coordinator.RoborockLocalClient.get_room_mapping", + return_value=[ + RoomMapping(16, "2362048"), + RoomMapping(17, "2362044"), + RoomMapping(18, "2362041"), + ], + ), + patch( + "homeassistant.components.roborock.coordinator.RoborockMqttClient.get_room_mapping", + return_value=[ + RoomMapping(16, "2362048"), + RoomMapping(17, "2362044"), + RoomMapping(18, "2362041"), + ], + ), ): yield diff --git a/tests/components/roborock/snapshots/test_vacuum.ambr b/tests/components/roborock/snapshots/test_vacuum.ambr new file mode 100644 index 00000000000..d03bec28125 --- /dev/null +++ b/tests/components/roborock/snapshots/test_vacuum.ambr @@ -0,0 +1,27 @@ +# serializer version: 1 +# name: test_get_maps + dict({ + 'vacuum.roborock_s7_maxv': dict({ + 'maps': list([ + dict({ + 'flag': 0, + 'name': 'Upstairs', + 'rooms': dict({ + 16: 'Example room 1', + 17: 'Example room 2', + 18: 'Example room 3', + }), + }), + dict({ + 'flag': 1, + 'name': 'Downstairs', + 'rooms': dict({ + 16: 'Example room 1', + 17: 'Example room 2', + 18: 'Example room 3', + }), + }), + ]), + }), + }) +# --- diff --git a/tests/components/roborock/test_vacuum.py b/tests/components/roborock/test_vacuum.py index a3d5854edd1..cc01acc29fd 100644 --- a/tests/components/roborock/test_vacuum.py +++ b/tests/components/roborock/test_vacuum.py @@ -7,8 +7,10 @@ from unittest.mock import patch import pytest from roborock import RoborockException from roborock.roborock_typing import RoborockCommand +from syrupy.assertion import SnapshotAssertion from homeassistant.components.roborock import DOMAIN +from homeassistant.components.roborock.const import GET_MAPS_SERVICE_NAME from homeassistant.components.vacuum import ( SERVICE_CLEAN_SPOT, SERVICE_LOCATE, @@ -154,3 +156,20 @@ async def test_failed_user_command( data, blocking=True, ) + + +async def test_get_maps( + hass: HomeAssistant, + bypass_api_fixture, + setup_entry: MockConfigEntry, + snapshot: SnapshotAssertion, +) -> None: + """Test that the service for maps correctly outputs rooms with the right name.""" + response = await hass.services.async_call( + DOMAIN, + GET_MAPS_SERVICE_NAME, + {ATTR_ENTITY_ID: ENTITY_ID}, + blocking=True, + return_response=True, + ) + assert response == snapshot