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

@@ -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()

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"

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")
+                    )

homeassistant/components/roborock/icons.json

@@ -105,5 +105,8 @@
         "default": "mdi:power-plug-off"
       }
     }
+  },
+  "services": {
+    "get_maps": "mdi:floor-plan"
   }
 }

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

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]

homeassistant/components/roborock/services.yaml

@@ -0,0 +1,4 @@
+get_maps:
+  target:
+    entity:
+      domain: vacuum

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."
+    }
   }
 }

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()]}

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
 

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',
+          }),
+        }),
+      ]),
+    }),
+  })
+# ---

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