Add ability to select current map for Roborock (home-assistant#120882)

Co-authored-by: J. Nick Koston <nick@koston.org>
jfparis · Jul 5, 2024 · ffc3958 · ffc3958
1 parent 45ab9ca
commit ffc3958
Show file tree

Hide file tree

Showing 11 changed files with 86 additions and 20 deletions.
diff --git a/homeassistant/components/roborock/binary_sensor.py b/homeassistant/components/roborock/binary_sensor.py
@@ -15,7 +15,6 @@
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.util import slugify
 
 from . import RoborockConfigEntry
 from .coordinator import RoborockDataUpdateCoordinator
@@ -97,7 +96,7 @@ def __init__(
     ) -> None:
         """Initialize the entity."""
         super().__init__(
-            f"{description.key}_{slugify(coordinator.duid)}",
+            f"{description.key}_{coordinator.duid_slug}",
             coordinator,
         )
         self.entity_description = description

diff --git a/homeassistant/components/roborock/button.py b/homeassistant/components/roborock/button.py
@@ -10,7 +10,6 @@
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.util import slugify
 
 from . import RoborockConfigEntry
 from .coordinator import RoborockDataUpdateCoordinator
@@ -90,7 +89,7 @@ def __init__(
     ) -> None:
         """Create a button entity."""
         super().__init__(
-            f"{entity_description.key}_{slugify(coordinator.duid)}",
+            f"{entity_description.key}_{coordinator.duid_slug}",
             coordinator.device_info,
             coordinator.api,
         )

diff --git a/homeassistant/components/roborock/coordinator.py b/homeassistant/components/roborock/coordinator.py
@@ -4,6 +4,7 @@
 
 import asyncio
 from datetime import timedelta
+from functools import cached_property
 import logging
 
 from roborock import HomeDataRoom
@@ -21,6 +22,7 @@
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+from homeassistant.util import slugify
 
 from .const import DOMAIN
 from .models import RoborockA01HassDeviceInfo, RoborockHassDeviceInfo, RoborockMapInfo
@@ -142,11 +144,16 @@ async def get_rooms(self) -> None:
                         self._home_data_rooms.get(room.iot_id, "Unknown")
                     )
 
-    @property
+    @cached_property
     def duid(self) -> str:
         """Get the unique id of the device as specified by Roborock."""
         return self.roborock_device_info.device.duid
 
+    @cached_property
+    def duid_slug(self) -> str:
+        """Get the slug of the duid."""
+        return slugify(self.duid)
+
 
 class RoborockDataUpdateCoordinatorA01(
     DataUpdateCoordinator[
@@ -191,7 +198,12 @@ async def release(self) -> None:
         """Disconnect from API."""
         await self.api.async_release()
 
-    @property
+    @cached_property
     def duid(self) -> str:
         """Get the unique id of the device as specified by Roborock."""
         return self.roborock_device_info.device.duid
+
+    @cached_property
+    def duid_slug(self) -> str:
+        """Get the slug of the duid."""
+        return slugify(self.duid)
diff --git a/homeassistant/components/roborock/number.py b/homeassistant/components/roborock/number.py
@@ -14,7 +14,6 @@
 from homeassistant.const import PERCENTAGE, EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.util import slugify
 
 from . import RoborockConfigEntry
 from .coordinator import RoborockDataUpdateCoordinator
@@ -77,7 +76,7 @@ async def async_setup_entry(
         else:
             valid_entities.append(
                 RoborockNumberEntity(
-                    f"{description.key}_{slugify(coordinator.duid)}",
+                    f"{description.key}_{coordinator.duid_slug}",
                     coordinator,
                     description,
                 )

diff --git a/homeassistant/components/roborock/select.py b/homeassistant/components/roborock/select.py
@@ -11,7 +11,6 @@
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.util import slugify
 
 from . import RoborockConfigEntry
 from .coordinator import RoborockDataUpdateCoordinator
@@ -79,6 +78,12 @@ async def async_setup_entry(
         )
         is not None
     )
+    async_add_entities(
+        RoborockCurrentMapSelectEntity(
+            f"selected_map_{coordinator.duid_slug}", coordinator
+        )
+        for coordinator in config_entry.runtime_data.v1
+    )
 
 
 class RoborockSelectEntity(RoborockCoordinatedEntityV1, SelectEntity):
@@ -95,7 +100,7 @@ def __init__(
         """Create a select entity."""
         self.entity_description = entity_description
         super().__init__(
-            f"{entity_description.key}_{slugify(coordinator.duid)}",
+            f"{entity_description.key}_{coordinator.duid_slug}",
             coordinator,
             entity_description.protocol_listener,
         )
@@ -112,3 +117,32 @@ async def async_select_option(self, option: str) -> None:
     def current_option(self) -> str | None:
         """Get the current status of the select entity from device_status."""
         return self.entity_description.value_fn(self._device_status)
+
+
+class RoborockCurrentMapSelectEntity(RoborockCoordinatedEntityV1, SelectEntity):
+    """A class to let you set the selected map on Roborock vacuum."""
+
+    _attr_entity_category = EntityCategory.DIAGNOSTIC
+    _attr_translation_key = "selected_map"
+
+    async def async_select_option(self, option: str) -> None:
+        """Set the option."""
+        for map_id, map_ in self.coordinator.maps.items():
+            if map_.name == option:
+                await self.send(
+                    RoborockCommand.LOAD_MULTI_MAP,
+                    [map_id],
+                )
+                break
+
+    @property
+    def options(self) -> list[str]:
+        """Gets all of the names of rooms that we are currently aware of."""
+        return [roborock_map.name for roborock_map in self.coordinator.maps.values()]
+
+    @property
+    def current_option(self) -> str | None:
+        """Get the current status of the select entity from device_status."""
+        if current_map := self.coordinator.current_map:
+            return self.coordinator.maps[current_map].name
+        return None
diff --git a/homeassistant/components/roborock/sensor.py b/homeassistant/components/roborock/sensor.py
@@ -30,7 +30,6 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
-from homeassistant.util import slugify
 
 from . import RoborockConfigEntry
 from .coordinator import RoborockDataUpdateCoordinator, RoborockDataUpdateCoordinatorA01
@@ -291,7 +290,7 @@ def __init__(
         """Initialize the entity."""
         self.entity_description = description
         super().__init__(
-            f"{description.key}_{slugify(coordinator.duid)}",
+            f"{description.key}_{coordinator.duid_slug}",
             coordinator,
             description.protocol_listener,
         )
@@ -316,7 +315,7 @@ def __init__(
     ) -> None:
         """Initialize the entity."""
         self.entity_description = description
-        super().__init__(f"{description.key}_{slugify(coordinator.duid)}", coordinator)
+        super().__init__(f"{description.key}_{coordinator.duid_slug}", coordinator)
 
     @property
     def native_value(self) -> StateType:

diff --git a/homeassistant/components/roborock/strings.json b/homeassistant/components/roborock/strings.json
@@ -298,6 +298,9 @@
           "custom": "[%key:component::roborock::entity::select::mop_mode::state::custom%]",
           "custom_water_flow": "Custom water flow"
         }
+      },
+      "selected_map": {
+        "name": "Selected map"
       }
     },
     "switch": {

diff --git a/homeassistant/components/roborock/switch.py b/homeassistant/components/roborock/switch.py
@@ -15,7 +15,6 @@
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.util import slugify
 
 from . import RoborockConfigEntry
 from .coordinator import RoborockDataUpdateCoordinator
@@ -125,7 +124,7 @@ async def async_setup_entry(
         else:
             valid_entities.append(
                 RoborockSwitch(
-                    f"{description.key}_{slugify(coordinator.duid)}",
+                    f"{description.key}_{coordinator.duid_slug}",
                     coordinator,
                     description,
                 )

diff --git a/homeassistant/components/roborock/time.py b/homeassistant/components/roborock/time.py
@@ -16,7 +16,6 @@
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.util import slugify
 
 from . import RoborockConfigEntry
 from .coordinator import RoborockDataUpdateCoordinator
@@ -141,7 +140,7 @@ async def async_setup_entry(
         else:
             valid_entities.append(
                 RoborockTimeEntity(
-                    f"{description.key}_{slugify(coordinator.duid)}",
+                    f"{description.key}_{coordinator.duid_slug}",
                     coordinator,
                     description,
                 )

diff --git a/homeassistant/components/roborock/vacuum.py b/homeassistant/components/roborock/vacuum.py
@@ -20,7 +20,6 @@
 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 . import RoborockConfigEntry
 from .const import DOMAIN, GET_MAPS_SERVICE_NAME
@@ -103,7 +102,7 @@ def __init__(
         StateVacuumEntity.__init__(self)
         RoborockCoordinatedEntityV1.__init__(
             self,
-            slugify(coordinator.duid),
+            coordinator.duid_slug,
             coordinator,
             listener_request=[
                 RoborockDataProtocol.FAN_POWER,

diff --git a/tests/components/roborock/test_select.py b/tests/components/roborock/test_select.py
@@ -1,13 +1,18 @@
 """Test Roborock Select platform."""
 
+import copy
 from unittest.mock import patch
 
 import pytest
 from roborock.exceptions import RoborockException
 
-from homeassistant.const import SERVICE_SELECT_OPTION
+from homeassistant.components.roborock import DOMAIN
+from homeassistant.const import SERVICE_SELECT_OPTION, STATE_UNKNOWN
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.setup import async_setup_component
+
+from .mock_data import PROP
 
 from tests.common import MockConfigEntry
 
@@ -17,6 +22,7 @@
     [
         ("select.roborock_s7_maxv_mop_mode", "deep"),
         ("select.roborock_s7_maxv_mop_intensity", "mild"),
+        ("select.roborock_s7_maxv_selected_map", "Downstairs"),
     ],
 )
 async def test_update_success(
@@ -62,3 +68,21 @@ async def test_update_failure(
             blocking=True,
             target={"entity_id": "select.roborock_s7_maxv_mop_mode"},
         )
+
+
+async def test_none_map_select(
+    hass: HomeAssistant,
+    bypass_api_fixture,
+    mock_roborock_entry: MockConfigEntry,
+) -> None:
+    """Test that the select entity correctly handles not having a current map."""
+    prop = copy.deepcopy(PROP)
+    # Set map status to None so that current map is never set
+    prop.status.map_status = None
+    with patch(
+        "homeassistant.components.roborock.coordinator.RoborockLocalClientV1.get_prop",
+        return_value=prop,
+    ):
+        await async_setup_component(hass, DOMAIN, {})
+    select_entity = hass.states.get("select.roborock_s7_maxv_selected_map")
+    assert select_entity.state == STATE_UNKNOWN