Add Roborock Q10 volume number entity

homeassistant/components/roborock/number.py

@@ -5,6 +5,8 @@
 import logging
 from typing import Any, override
 
+from roborock.devices.traits.b01 import Q10PropertiesApi
+from roborock.devices.traits.b01.q10 import SoundVolumeTrait
 from roborock.devices.traits.v1 import PropertiesApi
 from roborock.exceptions import RoborockException
 
@@ -17,11 +19,12 @@
 
 from .const import DOMAIN
 from .coordinator import (
+    RoborockB01Q10UpdateCoordinator,
     RoborockConfigEntry,
     RoborockCoordinatorType,
     RoborockDataUpdateCoordinator,
 )
-from .entity import RoborockEntityV1
+from .entity import RoborockCoordinatedEntityB01Q10, RoborockEntityV1
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -59,6 +62,37 @@ class RoborockNumberDescription(NumberEntityDescription):
 ]
 
 
+@dataclass(frozen=True, kw_only=True)
+class RoborockNumberDescriptionQ10(NumberEntityDescription):
+    """Class to describe a Roborock Q10 number entity."""
+
+    trait: Callable[[Q10PropertiesApi], SoundVolumeTrait | None]
+    """Function to get the trait backing the entity, if supported."""
+
+    get_value: Callable[[SoundVolumeTrait], float | None]
+    """Function to get the value from the trait."""
+
+    set_value: Callable[[SoundVolumeTrait, float], Coroutine[Any, Any, None]]
+    """Function to set the value on the trait."""
+
+
+Q10_NUMBER_DESCRIPTIONS: list[RoborockNumberDescriptionQ10] = [
+    RoborockNumberDescriptionQ10(
+        key="volume",
+        translation_key="volume",
+        native_min_value=0,
+        native_max_value=100,
+        native_unit_of_measurement=PERCENTAGE,
+        entity_category=EntityCategory.CONFIG,
+        trait=lambda api: api.volume,
+        get_value=lambda trait: (
+            float(trait.volume) if trait.volume is not None else None
+        ),
+        set_value=lambda trait, value: trait.set_volume(int(value)),
+    )
+]
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: RoborockConfigEntry,
@@ -72,18 +106,29 @@ def async_add_coordinator_entities(
         coordinator: RoborockCoordinatorType,
     ) -> None:
         """Add entities for a specific coordinator."""
-        if not isinstance(coordinator, RoborockDataUpdateCoordinator):
-            return
-        entities = [
-            RoborockNumberEntity(
-                f"{description.key}_{coordinator.duid_slug}",
-                coordinator=coordinator,
-                entity_description=description,
-                trait=trait,
+        entities: list[NumberEntity] = []
+        if isinstance(coordinator, RoborockDataUpdateCoordinator):
+            entities.extend(
+                RoborockNumberEntity(
+                    f"{description.key}_{coordinator.duid_slug}",
+                    coordinator=coordinator,
+                    entity_description=description,
+                    trait=trait,
+                )
+                for description in NUMBER_DESCRIPTIONS
+                if (trait := description.trait(coordinator.properties_api)) is not None
+            )
+        elif isinstance(coordinator, RoborockB01Q10UpdateCoordinator):
+            entities.extend(
+                RoborockNumberEntityQ10(
+                    f"{description.key}_{coordinator.duid_slug}",
+                    coordinator=coordinator,
+                    entity_description=description,
+                    trait=q10_trait,
+                )
+                for description in Q10_NUMBER_DESCRIPTIONS
+                if (q10_trait := description.trait(coordinator.api)) is not None
             )
-            for description in NUMBER_DESCRIPTIONS
-            if (trait := description.trait(coordinator.properties_api)) is not None
-        ]
         async_add_entities(entities)
 
     for coordinator in coordinators.values():
@@ -133,3 +178,45 @@ async def async_set_native_value(self, value: float) -> None:
                 translation_domain=DOMAIN,
                 translation_key="update_options_failed",
             ) from err
+
+
+class RoborockNumberEntityQ10(RoborockCoordinatedEntityB01Q10, NumberEntity):
+    """A class to set a numeric setting on a Roborock Q10 device."""
+
+    entity_description: RoborockNumberDescriptionQ10
+    coordinator: RoborockB01Q10UpdateCoordinator
+
+    def __init__(
+        self,
+        unique_id: str,
+        coordinator: RoborockB01Q10UpdateCoordinator,
+        entity_description: RoborockNumberDescriptionQ10,
+        trait: SoundVolumeTrait,
+    ) -> None:
+        """Create a number entity."""
+        self.entity_description = entity_description
+        self._trait = trait
+        super().__init__(unique_id, coordinator)
+
+    @override
+    async def async_added_to_hass(self) -> None:
+        """Register a trait listener for push-based state updates."""
+        await super().async_added_to_hass()
+        self.async_on_remove(self._trait.add_update_listener(self.async_write_ha_state))
+
+    @property
+    @override
+    def native_value(self) -> float | None:
+        """Get native value."""
+        return self.entity_description.get_value(self._trait)
+
+    @override
+    async def async_set_native_value(self, value: float) -> None:
+        """Set number value."""
+        try:
+            await self.entity_description.set_value(self._trait, value)
+        except RoborockException as err:
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="update_options_failed",
+            ) from err

tests/components/roborock/conftest.py

@@ -238,6 +238,15 @@ def create_b01_q10_trait() -> Mock:
     q10_trait.button_light.enable = AsyncMock()
     q10_trait.button_light.disable = AsyncMock()
 
+    q10_trait.volume = AsyncMock()
+    q10_trait.volume.volume = 50
+    volume_notify = attach_update_listeners(q10_trait.volume)
+
+    async def _set_volume(volume: int) -> None:
+        q10_trait.volume.volume = volume
+        volume_notify()
+
+    q10_trait.volume.set_volume = AsyncMock(side_effect=_set_volume)
     q10_trait.map = Mock()
     q10_trait.map.rooms = [
         Q10Room(id=9, raw_name="rr_bedroom", pixel_value=36, pixel_count=100),

tests/components/roborock/test_number.py

@@ -50,6 +50,74 @@ async def test_update_sound_volume(
     assert state.state == "3.0"
 
 
+async def test_q10_update_sound_volume(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    fake_q10_vacuum: FakeDevice,
+) -> None:
+    """Test changing the volume of a Q10 device."""
+    entity_id = "number.roborock_q10_s5_volume"
+
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "50.0"
+
+    await hass.services.async_call(
+        "number",
+        SERVICE_SET_VALUE,
+        service_data={ATTR_VALUE: 30.0},
+        blocking=True,
+        target={"entity_id": entity_id},
+    )
+
+    assert fake_q10_vacuum.b01_q10_properties is not None
+    fake_q10_vacuum.b01_q10_properties.volume.set_volume.assert_awaited_once_with(30)
+
+    # The trait listener pushes the new value into the entity state
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "30.0"
+
+
+async def test_q10_volume_unknown_value(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    fake_q10_vacuum: FakeDevice,
+) -> None:
+    """Test the Q10 entity reports unknown when the trait value is None."""
+    assert fake_q10_vacuum.b01_q10_properties is not None
+    fake_q10_vacuum.b01_q10_properties.volume.volume = None
+
+    await async_update_entity(hass, "number.roborock_q10_s5_volume")
+
+    state = hass.states.get("number.roborock_q10_s5_volume")
+    assert state is not None
+    assert state.state == STATE_UNKNOWN
+
+
+async def test_q10_volume_update_failed(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    fake_q10_vacuum: FakeDevice,
+) -> None:
+    """Test a failure while changing the volume of a Q10 device."""
+    assert fake_q10_vacuum.b01_q10_properties is not None
+    fake_q10_vacuum.b01_q10_properties.volume.set_volume.side_effect = RoborockTimeout
+
+    assert hass.states.get("number.roborock_q10_s5_volume") is not None
+
+    with pytest.raises(HomeAssistantError, match="Failed to update Roborock options"):
+        await hass.services.async_call(
+            "number",
+            SERVICE_SET_VALUE,
+            service_data={ATTR_VALUE: 30.0},
+            blocking=True,
+            target={"entity_id": "number.roborock_q10_s5_volume"},
+        )
+
+    fake_q10_vacuum.b01_q10_properties.volume.set_volume.assert_awaited_once_with(30)
+
+
 async def test_volume_unknown_value(
     hass: HomeAssistant,
     setup_entry: MockConfigEntry,