Add mopping abilities to Roborock (#91766)

Co-authored-by: Franck Nijhof <frenck@frenck.nl>

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]

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

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

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)

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

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