diff --git a/homeassistant/components/roborock/coordinator.py b/homeassistant/components/roborock/coordinator.py index 8e0a26d4916..6100d997d63 100644 --- a/homeassistant/components/roborock/coordinator.py +++ b/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( diff --git a/homeassistant/components/roborock/select.py b/homeassistant/components/roborock/select.py index bd376c03255..341dea0b267 100644 --- a/homeassistant/components/roborock/select.py +++ b/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): diff --git a/homeassistant/components/roborock/strings.json b/homeassistant/components/roborock/strings.json index fb360b1a784..0942dabea4c 100644 --- a/homeassistant/components/roborock/strings.json +++ b/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": { diff --git a/tests/components/roborock/conftest.py b/tests/components/roborock/conftest.py index 70f7ab0ad7a..7e3655782d4 100644 --- a/tests/components/roborock/conftest.py +++ b/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 diff --git a/tests/components/roborock/test_select.py b/tests/components/roborock/test_select.py index 8c05e13b1b0..95cc70d5612 100644 --- a/tests/components/roborock/test_select.py +++ b/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)