Add the ability to set Cleaning mode and mop mode for Q7 Vacs (#161725)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

homeassistant/components/roborock/coordinator.py

@@ -553,6 +553,8 @@ class RoborockB01Q7UpdateCoordinator(RoborockDataUpdateCoordinatorB01):
             RoborockB01Props.REAL_CLEAN_TIME,
             RoborockB01Props.HYPA,
             RoborockB01Props.WIND,
+            RoborockB01Props.WATER,
+            RoborockB01Props.MODE,
         ]
 
     async def _async_update_data(

homeassistant/components/roborock/select.py

@@ -1,10 +1,13 @@
 """Support for Roborock select."""
 
 import asyncio
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
+from typing import Any
 
-from roborock.data import RoborockDockDustCollectionModeCode
+from roborock import B01Props, CleanTypeMapping
+from roborock.data import RoborockDockDustCollectionModeCode, WaterLevelMapping
+from roborock.devices.traits.b01 import Q7PropertiesApi
 from roborock.devices.traits.v1 import PropertiesApi
 from roborock.devices.traits.v1.home import HomeTrait
 from roborock.devices.traits.v1.maps import MapsTrait
@@ -18,8 +21,12 @@ from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
 from .const import DOMAIN, MAP_SLEEP
-from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
-from .entity import RoborockCoordinatedEntityV1
+from .coordinator import (
+    RoborockB01Q7UpdateCoordinator,
+    RoborockConfigEntry,
+    RoborockDataUpdateCoordinator,
+)
+from .entity import RoborockCoordinatedEntityB01, RoborockCoordinatedEntityV1
 
 PARALLEL_UPDATES = 0
 
@@ -44,6 +51,42 @@ class RoborockSelectDescription(SelectEntityDescription):
     """Whether this entity is for the dock."""
 
 
+@dataclass(frozen=True, kw_only=True)
+class RoborockB01SelectDescription(SelectEntityDescription):
+    """Class to describe a Roborock B01 select entity."""
+
+    api_fn: Callable[[Q7PropertiesApi, str], Awaitable[Any]]
+    """Function to call the API."""
+
+    value_fn: Callable[[B01Props], str | None]
+    """Function to get the current value of the select entity."""
+
+    options_lambda: Callable[[Q7PropertiesApi], list[str] | None]
+    """Function to get all options of the select entity or returns None if not supported."""
+
+
+B01_SELECT_DESCRIPTIONS: list[RoborockB01SelectDescription] = [
+    RoborockB01SelectDescription(
+        key="water_flow",
+        translation_key="water_flow",
+        api_fn=lambda api, value: api.set_water_level(
+            WaterLevelMapping.from_value(value)
+        ),
+        value_fn=lambda data: data.water.value if data.water else None,
+        options_lambda=lambda _: [option.value for option in WaterLevelMapping],
+        entity_category=EntityCategory.CONFIG,
+    ),
+    RoborockB01SelectDescription(
+        key="cleaning_mode",
+        translation_key="cleaning_mode",
+        api_fn=lambda api, value: api.set_mode(CleanTypeMapping.from_value(value)),
+        value_fn=lambda data: data.mode.value if data.mode else None,
+        options_lambda=lambda _: list(CleanTypeMapping.keys()),
+        entity_category=EntityCategory.CONFIG,
+    ),
+]
+
+
 SELECT_DESCRIPTIONS: list[RoborockSelectDescription] = [
     RoborockSelectDescription(
         key="water_box_mode",
@@ -114,6 +157,52 @@ async def async_setup_entry(
         if (home_trait := coordinator.properties_api.home) is not None
         if (map_trait := coordinator.properties_api.maps) is not None
     )
+    async_add_entities(
+        RoborockB01SelectEntity(coordinator, description, options)
+        for coordinator in config_entry.runtime_data.b01
+        for description in B01_SELECT_DESCRIPTIONS
+        if isinstance(coordinator, RoborockB01Q7UpdateCoordinator)
+        if (options := description.options_lambda(coordinator.api)) is not None
+    )
+
+
+class RoborockB01SelectEntity(RoborockCoordinatedEntityB01, SelectEntity):
+    """Select entity for Roborock B01 devices."""
+
+    entity_description: RoborockB01SelectDescription
+    coordinator: RoborockB01Q7UpdateCoordinator
+
+    def __init__(
+        self,
+        coordinator: RoborockB01Q7UpdateCoordinator,
+        entity_description: RoborockB01SelectDescription,
+        options: list[str],
+    ) -> None:
+        """Initialize the entity."""
+        self.entity_description = entity_description
+        super().__init__(
+            f"{entity_description.key}_{coordinator.duid_slug}", coordinator
+        )
+        self._attr_options = options
+
+    async def async_select_option(self, option: str) -> None:
+        """Set the option."""
+        try:
+            await self.entity_description.api_fn(self.coordinator.api, option)
+        except RoborockException as err:
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="command_failed",
+                translation_placeholders={
+                    "command": self.entity_description.key,
+                },
+            ) from err
+        await self.coordinator.async_refresh()
+
+    @property
+    def current_option(self) -> str | None:
+        """Return the current option."""
+        return self.entity_description.value_fn(self.coordinator.data)
 
 
 class RoborockSelectEntity(RoborockCoordinatedEntityV1, SelectEntity):

homeassistant/components/roborock/strings.json

@@ -86,6 +86,14 @@
       }
     },
     "select": {
+      "cleaning_mode": {
+        "name": "Cleaning mode",
+        "state": {
+          "mop": "Mop only",
+          "vac_and_mop": "Vacuum and mop",
+          "vacuum": "Vacuum only"
+        }
+      },
       "dust_collection_mode": {
         "name": "Empty mode",
         "state": {
@@ -126,6 +134,14 @@
       },
       "selected_map": {
         "name": "Selected map"
+      },
+      "water_flow": {
+        "name": "Water flow",
+        "state": {
+          "high": "[%key:common::state::high%]",
+          "low": "[%key:common::state::low%]",
+          "medium": "[%key:common::state::medium%]"
+        }
       }
     },
     "sensor": {

tests/components/roborock/conftest.py

@@ -137,6 +137,8 @@ def create_b01_q7_trait() -> Mock:
     b01_trait.return_to_dock = AsyncMock(side_effect=return_to_dock_side_effect)
     b01_trait.find_me = AsyncMock()
     b01_trait.set_fan_speed = AsyncMock()
+    b01_trait.set_mode = AsyncMock()
+    b01_trait.set_water_level = AsyncMock()
     b01_trait.send = AsyncMock()
     return b01_trait
 

tests/components/roborock/test_select.py

@@ -4,8 +4,8 @@ from typing import Any
 from unittest.mock import AsyncMock, call
 
 import pytest
-from roborock import RoborockCommand
-from roborock.data.v1 import RoborockDockDustCollectionModeCode
+from roborock import CleanTypeMapping, RoborockCommand
+from roborock.data import RoborockDockDustCollectionModeCode, WaterLevelMapping
 from roborock.exceptions import RoborockException
 
 from homeassistant.components.roborock import DOMAIN
@@ -182,3 +182,99 @@ async def test_dust_collection_mode_none(
     select_entity = hass.states.get("select.roborock_s7_maxv_dock_empty_mode")
     assert select_entity
     assert select_entity.state == expected_state
+
+
+@pytest.fixture
+def q7_device(fake_devices: list[FakeDevice]) -> FakeDevice:
+    """Get the fake Q7 vacuum device."""
+    # The Q7 is the fourth device in the list (index 3) based on HOME_DATA
+    return fake_devices[3]
+
+
+async def test_update_success_q7_water_level(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    q7_device: FakeDevice,
+) -> None:
+    """Test allowed changing values for Q7 water flow select entity."""
+    entity_id = "select.roborock_q7_water_flow"
+    assert hass.states.get(entity_id) is not None
+
+    # Test setting value
+    await hass.services.async_call(
+        "select",
+        SERVICE_SELECT_OPTION,
+        service_data={"option": "high"},
+        blocking=True,
+        target={"entity_id": entity_id},
+    )
+
+    assert q7_device.b01_q7_properties
+    assert q7_device.b01_q7_properties.set_water_level.call_count == 1
+    q7_device.b01_q7_properties.set_water_level.assert_called_with(
+        WaterLevelMapping.HIGH
+    )
+
+
+async def test_update_failure_q7_water_level(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    q7_device: FakeDevice,
+) -> None:
+    """Test failure when setting Q7 water flow."""
+    assert q7_device.b01_q7_properties
+    q7_device.b01_q7_properties.set_water_level.side_effect = RoborockException
+    entity_id = "select.roborock_q7_water_flow"
+    assert hass.states.get(entity_id) is not None
+
+    with pytest.raises(HomeAssistantError, match="Error while calling water_flow"):
+        await hass.services.async_call(
+            "select",
+            SERVICE_SELECT_OPTION,
+            service_data={"option": "high"},
+            blocking=True,
+            target={"entity_id": entity_id},
+        )
+
+
+async def test_update_failure_q7_cleaning_mode(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    q7_device: FakeDevice,
+) -> None:
+    """Test failure when setting Q7 cleaning mode."""
+    assert q7_device.b01_q7_properties
+    q7_device.b01_q7_properties.set_mode.side_effect = RoborockException
+
+    with pytest.raises(HomeAssistantError, match="Error while calling cleaning_mode"):
+        await hass.services.async_call(
+            "select",
+            SERVICE_SELECT_OPTION,
+            service_data={"option": "vacuum"},
+            blocking=True,
+            target={"entity_id": "select.roborock_q7_cleaning_mode"},
+        )
+
+
+async def test_update_success_q7_cleaning_mode(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    q7_device: FakeDevice,
+) -> None:
+    """Test allowed changing values for Q7 cleaning mode select entity."""
+    entity_id = "select.roborock_q7_cleaning_mode"
+    assert hass.states.get(entity_id) is not None
+
+    # Test setting value
+    await hass.services.async_call(
+        "select",
+        SERVICE_SELECT_OPTION,
+        service_data={"option": "vacuum"},
+        blocking=True,
+        target={"entity_id": entity_id},
+    )
+
+    assert q7_device.b01_q7_properties
+    assert q7_device.b01_q7_properties.set_mode.call_count == 1
+
+    q7_device.b01_q7_properties.set_mode.assert_called_with(CleanTypeMapping.VACUUM)