From b5bf88215c271077c3b3ee1faa996ee60f7c5c6c Mon Sep 17 00:00:00 2001 From: Luke Date: Mon, 24 Apr 2023 16:53:45 -0400 Subject: [PATCH] Add mopping abilities to Roborock (#91766) Co-authored-by: Franck Nijhof --- homeassistant/components/roborock/const.py | 2 +- .../components/roborock/coordinator.py | 2 +- homeassistant/components/roborock/device.py | 6 +- homeassistant/components/roborock/select.py | 116 ++++++++++++++++++ .../components/roborock/strings.json | 23 ++++ tests/components/roborock/test_select.py | 58 +++++++++ 6 files changed, 204 insertions(+), 3 deletions(-) create mode 100644 homeassistant/components/roborock/select.py create mode 100644 tests/components/roborock/test_select.py diff --git a/homeassistant/components/roborock/const.py b/homeassistant/components/roborock/const.py index 597d1923f20..61a9a70dd20 100644 --- a/homeassistant/components/roborock/const.py +++ b/homeassistant/components/roborock/const.py @@ -6,4 +6,4 @@ CONF_ENTRY_CODE = "code" CONF_BASE_URL = "base_url" CONF_USER_DATA = "user_data" -PLATFORMS = [Platform.VACUUM] +PLATFORMS = [Platform.VACUUM, Platform.SELECT] diff --git a/homeassistant/components/roborock/coordinator.py b/homeassistant/components/roborock/coordinator.py index daf3257d26a..997c0a6acb8 100644 --- a/homeassistant/components/roborock/coordinator.py +++ b/homeassistant/components/roborock/coordinator.py @@ -74,7 +74,7 @@ class RoborockDataUpdateCoordinator( async def _async_update_data(self) -> dict[str, RoborockDeviceProp]: """Update data via library.""" try: - asyncio.gather( + await asyncio.gather( *( self._update_device_prop(device_info) for device_info in self.devices_info.values() diff --git a/homeassistant/components/roborock/device.py b/homeassistant/components/roborock/device.py index 86dcbcade01..e544147e9b8 100644 --- a/homeassistant/components/roborock/device.py +++ b/homeassistant/components/roborock/device.py @@ -59,4 +59,8 @@ class RoborockCoordinatedEntity(CoordinatorEntity[RoborockDataUpdateCoordinator] self, command: RoborockCommand, params: dict[str, Any] | list[Any] | None = None ) -> dict: """Send a command to a vacuum cleaner.""" - return await self.coordinator.api.send_command(self._device_id, command, params) + response = await self.coordinator.api.send_command( + self._device_id, command, params + ) + await self.coordinator.async_request_refresh() + return response diff --git a/homeassistant/components/roborock/select.py b/homeassistant/components/roborock/select.py new file mode 100644 index 00000000000..646c4904854 --- /dev/null +++ b/homeassistant/components/roborock/select.py @@ -0,0 +1,116 @@ +"""Support for Roborock select.""" +from collections.abc import Callable +from dataclasses import dataclass + +from roborock.code_mappings import RoborockMopIntensityCode, RoborockMopModeCode +from roborock.containers import Status +from roborock.exceptions import RoborockException +from roborock.typing import RoborockCommand + +from homeassistant.components.select import SelectEntity, SelectEntityDescription +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import HomeAssistant +from homeassistant.exceptions import HomeAssistantError +from homeassistant.helpers.entity_platform import AddEntitiesCallback +from homeassistant.util import slugify + +from .const import DOMAIN +from .coordinator import RoborockDataUpdateCoordinator +from .device import RoborockCoordinatedEntity +from .models import RoborockHassDeviceInfo + + +@dataclass +class RoborockSelectDescriptionMixin: + """Define an entity description mixin for select entities.""" + + api_command: RoborockCommand + value_fn: Callable[[Status], str] + options_lambda: Callable[[str], list[int]] + + +@dataclass +class RoborockSelectDescription( + SelectEntityDescription, RoborockSelectDescriptionMixin +): + """Class to describe an Roborock select entity.""" + + +SELECT_DESCRIPTIONS: list[RoborockSelectDescription] = [ + RoborockSelectDescription( + key="water_box_mode", + translation_key="mop_intensity", + options=RoborockMopIntensityCode.values(), + api_command=RoborockCommand.SET_WATER_BOX_CUSTOM_MODE, + value_fn=lambda data: data.water_box_mode, + options_lambda=lambda data: [ + k for k, v in RoborockMopIntensityCode.items() if v == data + ], + ), + RoborockSelectDescription( + key="mop_mode", + translation_key="mop_mode", + options=RoborockMopModeCode.values(), + api_command=RoborockCommand.SET_MOP_MODE, + value_fn=lambda data: data.mop_mode, + options_lambda=lambda data: [ + k for k, v in RoborockMopModeCode.items() if v == data + ], + ), +] + + +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +) -> None: + """Set up Roborock select platform.""" + + coordinator: RoborockDataUpdateCoordinator = hass.data[DOMAIN][ + config_entry.entry_id + ] + async_add_entities( + RoborockSelectEntity( + f"{description.key}_{slugify(device_id)}", + device_info, + coordinator, + description, + ) + for device_id, device_info in coordinator.devices_info.items() + for description in SELECT_DESCRIPTIONS + ) + + +class RoborockSelectEntity(RoborockCoordinatedEntity, SelectEntity): + """A class to let you set options on a Roborock vacuum where the potential options are fixed.""" + + entity_description: RoborockSelectDescription + + def __init__( + self, + unique_id: str, + device_info: RoborockHassDeviceInfo, + coordinator: RoborockDataUpdateCoordinator, + entity_description: RoborockSelectDescription, + ) -> None: + """Create a select entity.""" + self.entity_description = entity_description + super().__init__(unique_id, device_info, coordinator) + + async def async_select_option(self, option: str) -> None: + """Set the mop intensity.""" + try: + await self.send( + self.entity_description.api_command, + self.entity_description.options_lambda(option), + ) + except RoborockException as err: + raise HomeAssistantError( + f"Error while setting {self.entity_description.key} to {option}" + ) from err + + @property + 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) diff --git a/homeassistant/components/roborock/strings.json b/homeassistant/components/roborock/strings.json index eb4a5e7d901..7e755a0c41f 100644 --- a/homeassistant/components/roborock/strings.json +++ b/homeassistant/components/roborock/strings.json @@ -22,5 +22,28 @@ "abort": { "already_configured": "[%key:common::config_flow::abort::already_configured_device%]" } + }, + "entity": { + "select": { + "mop_mode": { + "name": "Mop mode", + "state": { + "standard": "Standard", + "deep": "Deep", + "deep_plus": "Deep+", + "custom": "Custom" + } + }, + "mop_intensity": { + "name": "Mop intensity", + "state": { + "off": "Off", + "mild": "Mild", + "moderate": "Moderate", + "intense": "Intense", + "custom": "Custom" + } + } + } } } diff --git a/tests/components/roborock/test_select.py b/tests/components/roborock/test_select.py new file mode 100644 index 00000000000..3b0ba8183b3 --- /dev/null +++ b/tests/components/roborock/test_select.py @@ -0,0 +1,58 @@ +"""Test Roborock Select platform.""" +from unittest.mock import patch + +import pytest +from roborock.exceptions import RoborockException + +from homeassistant.const import SERVICE_SELECT_OPTION +from homeassistant.core import HomeAssistant +from homeassistant.exceptions import HomeAssistantError + +from tests.common import MockConfigEntry + + +@pytest.mark.parametrize( + ("entity_id", "value"), + [ + ("select.roborock_s7_maxv_mop_mode", "deep"), + ("select.roborock_s7_maxv_mop_intensity", "mild"), + ], +) +async def test_update_success( + hass: HomeAssistant, + bypass_api_fixture, + setup_entry: MockConfigEntry, + entity_id: str, + value: str, +) -> None: + """Test allowed changing values for select entities.""" + with patch( + "homeassistant.components.roborock.coordinator.RoborockLocalClient.send_message" + ) as mock_send_message: + await hass.services.async_call( + "select", + SERVICE_SELECT_OPTION, + service_data={"option": value}, + blocking=True, + target={"entity_id": entity_id}, + ) + assert mock_send_message.assert_called_once + + +async def test_update_failure( + hass: HomeAssistant, + bypass_api_fixture, + setup_entry: MockConfigEntry, +) -> None: + """Test that changing a value will raise a homeassistanterror when it fails.""" + with patch( + "homeassistant.components.roborock.coordinator.RoborockLocalClient.send_message", + side_effect=RoborockException(), + ), pytest.raises(HomeAssistantError): + await hass.services.async_call( + "select", + SERVICE_SELECT_OPTION, + service_data={"option": "deep"}, + blocking=True, + target={"entity_id": "select.roborock_s7_maxv_mop_mode"}, + )