Add select entities for Roborock q10 s5+

homeassistant/components/roborock/select.py

@@ -20,6 +20,7 @@
     ZeoSpin,
     ZeoTemperature,
 )
+from roborock.data.b01_q10.b01_q10_code_mappings import YXCleanType
 from roborock.devices.traits.b01 import Q7PropertiesApi
 from roborock.devices.traits.v1 import PropertiesApi
 from roborock.devices.traits.v1.home import HomeTrait
@@ -37,13 +38,15 @@
 from .const import DOMAIN, MAP_SLEEP
 from .coordinator import (
     RoborockB01Q7UpdateCoordinator,
+    RoborockB01Q10UpdateCoordinator,
     RoborockConfigEntry,
     RoborockDataUpdateCoordinator,
     RoborockDataUpdateCoordinatorA01,
 )
 from .entity import (
     RoborockCoordinatedEntityA01,
     RoborockCoordinatedEntityB01Q7,
+    RoborockCoordinatedEntityB01Q10,
     RoborockCoordinatedEntityV1,
 )
 
@@ -266,6 +269,10 @@ async def async_setup_entry(
         for description in A01_SELECT_DESCRIPTIONS
         if description.data_protocol in coordinator.request_protocols
     )
+    async_add_entities(
+        RoborockQ10CleanModeSelectEntity(coordinator)
+        for coordinator in config_entry.runtime_data.b01_q10
+    )
 
 
 class RoborockB01SelectEntity(RoborockCoordinatedEntityB01Q7, SelectEntity):
@@ -466,3 +473,59 @@ def current_option(self) -> str | None:
             self.entity_description.key,
         )
         return str(current_value)
+
+
+class RoborockQ10CleanModeSelectEntity(RoborockCoordinatedEntityB01Q10, SelectEntity):
+    """Select entity for Q10 cleaning mode."""
+
+    _attr_entity_category = EntityCategory.CONFIG
+    _attr_translation_key = "cleaning_mode"
+    coordinator: RoborockB01Q10UpdateCoordinator
+
+    def __init__(
+        self,
+        coordinator: RoborockB01Q10UpdateCoordinator,
+    ) -> None:
+        """Create a select entity for Q10 cleaning mode."""
+        super().__init__(
+            f"cleaning_mode_{coordinator.duid_slug}",
+            coordinator,
+        )
+
+    async def async_added_to_hass(self) -> None:
+        """Register trait listener for push-based status updates."""
+        await super().async_added_to_hass()
+        self.async_on_remove(
+            self.coordinator.api.status.add_update_listener(self.async_write_ha_state)
+        )
+
+    @property
+    def options(self) -> list[str]:
+        """Return available cleaning modes."""
+        return [mode.value for mode in YXCleanType if mode != YXCleanType.UNKNOWN]
+
+    @property
+    def current_option(self) -> str | None:
+        """Get the current cleaning mode."""
+        clean_mode = self.coordinator.api.status.clean_mode
+        if clean_mode is None or clean_mode == YXCleanType.UNKNOWN:
+            return None
+        return clean_mode.value
+
+    async def async_select_option(self, option: str) -> None:
+        """Set the cleaning mode."""
+        try:
+            mode = YXCleanType.from_value(option)
+        except ValueError as err:
+            raise ServiceValidationError(
+                translation_domain=DOMAIN,
+                translation_key="select_option_failed",
+            ) from err
+        try:
+            await self.coordinator.api.vacuum.set_clean_mode(mode)
+        except RoborockException as err:
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="command_failed",
+                translation_placeholders={"command": "cleaning_mode"},
+            ) from err

tests/components/roborock/test_select.py

@@ -10,6 +10,7 @@
     WaterLevelMapping,
     ZeoProgram,
 )
+from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP, YXCleanType
 from roborock.exceptions import RoborockException
 from roborock.roborock_message import RoborockZeoProtocol
 
@@ -383,3 +384,98 @@ async def test_update_failure_zeo_invalid_option() -> None:
         await entity.async_select_option("invalid_option")
 
     coordinator.api.set_value.assert_not_called()
+
+
+async def test_q10_cleaning_mode_select_current_option(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    fake_q10_vacuum: FakeDevice,
+) -> None:
+    """Test Q10 cleaning mode select entity current option."""
+    entity_id = "select.roborock_q10_s5_cleaning_mode"
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == STATE_UNKNOWN
+    options = state.attributes.get("options")
+    assert options is not None
+    assert set(options) == {"vac_and_mop", "vacuum", "mop"}
+
+    assert fake_q10_vacuum.b01_q10_properties
+    fake_q10_vacuum.b01_q10_properties.status.update_from_dps(
+        {B01_Q10_DP.CLEAN_MODE: YXCleanType.VAC_AND_MOP.code}
+    )
+    await hass.async_block_till_done()
+
+    updated_state = hass.states.get(entity_id)
+    assert updated_state is not None
+    assert updated_state.state == "vac_and_mop"
+
+
+async def test_q10_cleaning_mode_select_update_success(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    fake_q10_vacuum: FakeDevice,
+) -> None:
+    """Test setting Q10 cleaning mode select option."""
+    entity_id = "select.roborock_q10_s5_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": "vac_and_mop"},
+        blocking=True,
+        target={"entity_id": entity_id},
+    )
+
+    assert fake_q10_vacuum.b01_q10_properties
+    assert fake_q10_vacuum.b01_q10_properties.vacuum.set_clean_mode.call_count == 1
+    fake_q10_vacuum.b01_q10_properties.vacuum.set_clean_mode.assert_called_once_with(
+        YXCleanType.VAC_AND_MOP
+    )
+
+
+async def test_q10_cleaning_mode_select_update_failure(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    fake_q10_vacuum: FakeDevice,
+) -> None:
+    """Test failure when setting Q10 cleaning mode."""
+    assert fake_q10_vacuum.b01_q10_properties
+    fake_q10_vacuum.b01_q10_properties.vacuum.set_clean_mode.side_effect = (
+        RoborockException
+    )
+    entity_id = "select.roborock_q10_s5_cleaning_mode"
+    assert hass.states.get(entity_id) is not None
+
+    with pytest.raises(HomeAssistantError, match="cleaning_mode"):
+        await hass.services.async_call(
+            "select",
+            SERVICE_SELECT_OPTION,
+            service_data={"option": "vac_and_mop"},
+            blocking=True,
+            target={"entity_id": entity_id},
+        )
+
+
+async def test_q10_cleaning_mode_select_invalid_option(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    fake_q10_vacuum: FakeDevice,
+) -> None:
+    """Test that an invalid option raises ServiceValidationError and does not call set_clean_mode."""
+    entity_id = "select.roborock_q10_s5_cleaning_mode"
+    assert hass.states.get(entity_id) is not None
+
+    with pytest.raises(ServiceValidationError):
+        await hass.services.async_call(
+            "select",
+            SERVICE_SELECT_OPTION,
+            service_data={"option": "invalid_option"},
+            blocking=True,
+            target={"entity_id": entity_id},
+        )
+
+    assert fake_q10_vacuum.b01_q10_properties
+    fake_q10_vacuum.b01_q10_properties.vacuum.set_clean_mode.assert_not_called()