From ffc39585ed9ef769f87d2bb448850a7c280093e7 Mon Sep 17 00:00:00 2001 From: Luke Lashley Date: Fri, 5 Jul 2024 10:38:26 -0400 Subject: [PATCH] Add ability to select current map for Roborock (#120882) Co-authored-by: J. Nick Koston --- .../components/roborock/binary_sensor.py | 3 +- homeassistant/components/roborock/button.py | 3 +- .../components/roborock/coordinator.py | 16 +++++++- homeassistant/components/roborock/number.py | 3 +- homeassistant/components/roborock/select.py | 38 ++++++++++++++++++- homeassistant/components/roborock/sensor.py | 5 +-- .../components/roborock/strings.json | 3 ++ homeassistant/components/roborock/switch.py | 3 +- homeassistant/components/roborock/time.py | 3 +- homeassistant/components/roborock/vacuum.py | 3 +- tests/components/roborock/test_select.py | 26 ++++++++++++- 11 files changed, 86 insertions(+), 20 deletions(-) diff --git a/homeassistant/components/roborock/binary_sensor.py b/homeassistant/components/roborock/binary_sensor.py index 779d3ee234d..fb35a50c210 100644 --- a/homeassistant/components/roborock/binary_sensor.py +++ b/homeassistant/components/roborock/binary_sensor.py @@ -15,7 +15,6 @@ from homeassistant.components.binary_sensor import ( 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 @@ class RoborockBinarySensorEntity(RoborockCoordinatedEntityV1, BinarySensorEntity ) -> 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 index 50d84e37a44..31421320c41 100644 --- a/homeassistant/components/roborock/button.py +++ b/homeassistant/components/roborock/button.py @@ -10,7 +10,6 @@ from homeassistant.components.button import ButtonEntity, ButtonEntityDescriptio 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 @@ class RoborockButtonEntity(RoborockEntityV1, ButtonEntity): ) -> 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 index 430e2815a7b..a0e441201bb 100644 --- a/homeassistant/components/roborock/coordinator.py +++ b/homeassistant/components/roborock/coordinator.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio from datetime import timedelta +from functools import cached_property import logging from roborock import HomeDataRoom @@ -21,6 +22,7 @@ from homeassistant.helpers import device_registry as dr 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 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]): 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 @@ class RoborockDataUpdateCoordinatorA01( """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 index e86f07ad204..92552ca85d8 100644 --- a/homeassistant/components/roborock/number.py +++ b/homeassistant/components/roborock/number.py @@ -14,7 +14,6 @@ from homeassistant.components.number import NumberEntity, NumberEntityDescriptio 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 index 8966652c24d..f047ec475c2 100644 --- a/homeassistant/components/roborock/select.py +++ b/homeassistant/components/roborock/select.py @@ -11,7 +11,6 @@ from homeassistant.components.select import SelectEntity, SelectEntityDescriptio 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 @@ class RoborockSelectEntity(RoborockCoordinatedEntityV1, SelectEntity): """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 @@ class RoborockSelectEntity(RoborockCoordinatedEntityV1, SelectEntity): 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 index 71c996f0b53..36ee5fb02ce 100644 --- a/homeassistant/components/roborock/sensor.py +++ b/homeassistant/components/roborock/sensor.py @@ -30,7 +30,6 @@ from homeassistant.const import ( 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 @@ class RoborockSensorEntity(RoborockCoordinatedEntityV1, SensorEntity): """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 @@ class RoborockSensorEntityA01(RoborockCoordinatedEntityA01, SensorEntity): ) -> 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 index c7fc34386fd..362e0e4aff8 100644 --- 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 index 6cc562fb533..ef46fe61415 100644 --- a/homeassistant/components/roborock/switch.py +++ b/homeassistant/components/roborock/switch.py @@ -15,7 +15,6 @@ from homeassistant.components.switch import SwitchEntity, SwitchEntityDescriptio 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 index b0fbb18ed56..1136170192d 100644 --- a/homeassistant/components/roborock/time.py +++ b/homeassistant/components/roborock/time.py @@ -16,7 +16,6 @@ from homeassistant.components.time import TimeEntity, TimeEntityDescription 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 index 90f5002a23e..f7fc58161a8 100644 --- a/homeassistant/components/roborock/vacuum.py +++ b/homeassistant/components/roborock/vacuum.py @@ -20,7 +20,6 @@ from homeassistant.components.vacuum import ( 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 @@ class RoborockVacuum(RoborockCoordinatedEntityV1, StateVacuumEntity): 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 index c8626818749..ce846107d93 100644 --- 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 @@ from tests.common import MockConfigEntry [ ("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