Add listeners for roborock (#103651)

* Add listeners for roborock

* add tests

* decrease test complexity

homeassistant/components/roborock/device.py

@@ -5,7 +5,7 @@ from typing import Any
 from roborock.api import AttributeCache, RoborockClient
 from roborock.cloud_api import RoborockMqttClient
 from roborock.command_cache import CacheableAttribute
-from roborock.containers import Status
+from roborock.containers import Consumable, Status
 from roborock.exceptions import RoborockException
 from roborock.roborock_typing import RoborockCommand
 
@@ -97,3 +97,12 @@ class RoborockCoordinatedEntity(
         res = await super().send(command, params)
         await self.coordinator.async_refresh()
         return res
+
+    def _update_from_listener(self, value: Status | Consumable):
+        """Update the status or consumable data from a listener and then write the new entity state."""
+        if isinstance(value, Status):
+            self.coordinator.roborock_device_info.props.status = value
+        else:
+            self.coordinator.roborock_device_info.props.consumable = value
+        self.coordinator.data = self.coordinator.roborock_device_info.props
+        self.async_write_ha_state()

homeassistant/components/roborock/select.py

@@ -3,6 +3,7 @@ from collections.abc import Callable
 from dataclasses import dataclass
 
 from roborock.containers import Status
+from roborock.roborock_message import RoborockDataProtocol
 from roborock.roborock_typing import RoborockCommand
 
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
@@ -37,6 +38,8 @@ class RoborockSelectDescription(
 ):
     """Class to describe an Roborock select entity."""
 
+    protocol_listener: RoborockDataProtocol | None = None
+
 
 SELECT_DESCRIPTIONS: list[RoborockSelectDescription] = [
     RoborockSelectDescription(
@@ -49,6 +52,7 @@ SELECT_DESCRIPTIONS: list[RoborockSelectDescription] = [
         if data.water_box_mode is not None
         else None,
         parameter_lambda=lambda key, status: [status.get_mop_intensity_code(key)],
+        protocol_listener=RoborockDataProtocol.WATER_BOX_MODE,
     ),
     RoborockSelectDescription(
         key="mop_mode",
@@ -105,6 +109,8 @@ class RoborockSelectEntity(RoborockCoordinatedEntity, SelectEntity):
         self.entity_description = entity_description
         super().__init__(unique_id, coordinator)
         self._attr_options = options
+        if (protocol := self.entity_description.protocol_listener) is not None:
+            self.api.add_listener(protocol, self._update_from_listener, self.api.cache)
 
     async def async_select_option(self, option: str) -> None:
         """Set the option."""

homeassistant/components/roborock/sensor.py

@@ -11,6 +11,7 @@ from roborock.containers import (
     RoborockErrorCode,
     RoborockStateCode,
 )
+from roborock.roborock_message import RoborockDataProtocol
 from roborock.roborock_typing import DeviceProp
 
 from homeassistant.components.sensor import (
@@ -48,6 +49,8 @@ class RoborockSensorDescription(
 ):
     """A class that describes Roborock sensors."""
 
+    protocol_listener: RoborockDataProtocol | None = None
+
 
 def _dock_error_value_fn(properties: DeviceProp) -> str | None:
     if (
@@ -67,6 +70,7 @@ SENSOR_DESCRIPTIONS = [
         translation_key="main_brush_time_left",
         value_fn=lambda data: data.consumable.main_brush_time_left,
         entity_category=EntityCategory.DIAGNOSTIC,
+        protocol_listener=RoborockDataProtocol.MAIN_BRUSH_WORK_TIME,
     ),
     RoborockSensorDescription(
         native_unit_of_measurement=UnitOfTime.SECONDS,
@@ -76,6 +80,7 @@ SENSOR_DESCRIPTIONS = [
         translation_key="side_brush_time_left",
         value_fn=lambda data: data.consumable.side_brush_time_left,
         entity_category=EntityCategory.DIAGNOSTIC,
+        protocol_listener=RoborockDataProtocol.SIDE_BRUSH_WORK_TIME,
     ),
     RoborockSensorDescription(
         native_unit_of_measurement=UnitOfTime.SECONDS,
@@ -85,6 +90,7 @@ SENSOR_DESCRIPTIONS = [
         translation_key="filter_time_left",
         value_fn=lambda data: data.consumable.filter_time_left,
         entity_category=EntityCategory.DIAGNOSTIC,
+        protocol_listener=RoborockDataProtocol.FILTER_WORK_TIME,
     ),
     RoborockSensorDescription(
         native_unit_of_measurement=UnitOfTime.SECONDS,
@@ -120,6 +126,7 @@ SENSOR_DESCRIPTIONS = [
         value_fn=lambda data: data.status.state_name,
         entity_category=EntityCategory.DIAGNOSTIC,
         options=RoborockStateCode.keys(),
+        protocol_listener=RoborockDataProtocol.STATE,
     ),
     RoborockSensorDescription(
         key="cleaning_area",
@@ -145,6 +152,7 @@ SENSOR_DESCRIPTIONS = [
         value_fn=lambda data: data.status.error_code_name,
         entity_category=EntityCategory.DIAGNOSTIC,
         options=RoborockErrorCode.keys(),
+        protocol_listener=RoborockDataProtocol.ERROR_CODE,
     ),
     RoborockSensorDescription(
         key="battery",
@@ -152,6 +160,7 @@ SENSOR_DESCRIPTIONS = [
         entity_category=EntityCategory.DIAGNOSTIC,
         native_unit_of_measurement=PERCENTAGE,
         device_class=SensorDeviceClass.BATTERY,
+        protocol_listener=RoborockDataProtocol.BATTERY,
     ),
     RoborockSensorDescription(
         key="last_clean_start",
@@ -238,6 +247,8 @@ class RoborockSensorEntity(RoborockCoordinatedEntity, SensorEntity):
         """Initialize the entity."""
         super().__init__(unique_id, coordinator)
         self.entity_description = description
+        if (protocol := self.entity_description.protocol_listener) is not None:
+            self.api.add_listener(protocol, self._update_from_listener, self.api.cache)
 
     @property
     def native_value(self) -> StateType | datetime.datetime:

homeassistant/components/roborock/vacuum.py

@@ -2,6 +2,7 @@
 from typing import Any
 
 from roborock.code_mappings import RoborockStateCode
+from roborock.roborock_message import RoborockDataProtocol
 from roborock.roborock_typing import RoborockCommand
 
 from homeassistant.components.vacuum import (
@@ -94,6 +95,12 @@ class RoborockVacuum(RoborockCoordinatedEntity, StateVacuumEntity):
         StateVacuumEntity.__init__(self)
         RoborockCoordinatedEntity.__init__(self, unique_id, coordinator)
         self._attr_fan_speed_list = self._device_status.fan_power_options
+        self.api.add_listener(
+            RoborockDataProtocol.FAN_POWER, self._update_from_listener, self.api.cache
+        )
+        self.api.add_listener(
+            RoborockDataProtocol.STATE, self._update_from_listener, self.api.cache
+        )
 
     @property
     def state(self) -> str | None:

tests/components/roborock/test_sensor.py

@@ -1,14 +1,20 @@
 """Test Roborock Sensors."""
+from unittest.mock import patch
 
+from roborock import DeviceData, HomeDataDevice
+from roborock.cloud_api import RoborockMqttClient
 from roborock.const import (
     FILTER_REPLACE_TIME,
     MAIN_BRUSH_REPLACE_TIME,
     SENSOR_DIRTY_REPLACE_TIME,
     SIDE_BRUSH_REPLACE_TIME,
 )
+from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
 
 from homeassistant.core import HomeAssistant
 
+from .mock_data import CONSUMABLE, STATUS, USER_DATA
+
 from tests.common import MockConfigEntry
 
 
@@ -47,3 +53,41 @@ async def test_sensors(hass: HomeAssistant, setup_entry: MockConfigEntry) -> Non
         hass.states.get("sensor.roborock_s7_maxv_last_clean_end").state
         == "2023-01-01T03:43:58+00:00"
     )
+
+
+async def test_listener_update(
+    hass: HomeAssistant, setup_entry: MockConfigEntry
+) -> None:
+    """Test that when we receive a mqtt topic, we successfully update the entity."""
+    assert hass.states.get("sensor.roborock_s7_maxv_status").state == "charging"
+    # Listeners are global based on uuid - so this is okay
+    client = RoborockMqttClient(
+        USER_DATA, DeviceData(device=HomeDataDevice("abc123", "", "", "", ""), model="")
+    )
+    # Test Status
+    with patch("roborock.api.AttributeCache.value", STATUS.as_dict()):
+        # Symbolizes a mqtt message coming in
+        client.on_message_received(
+            [
+                RoborockMessage(
+                    protocol=RoborockMessageProtocol.GENERAL_REQUEST,
+                    payload=b'{"t": 1699464794, "dps": {"121": 5}}',
+                )
+            ]
+        )
+    # Test consumable
+    assert hass.states.get("sensor.roborock_s7_maxv_filter_time_left").state == str(
+        FILTER_REPLACE_TIME - 74382
+    )
+    with patch("roborock.api.AttributeCache.value", CONSUMABLE.as_dict()):
+        client.on_message_received(
+            [
+                RoborockMessage(
+                    protocol=RoborockMessageProtocol.GENERAL_REQUEST,
+                    payload=b'{"t": 1699464794, "dps": {"127": 743}}',
+                )
+            ]
+        )
+    assert hass.states.get("sensor.roborock_s7_maxv_filter_time_left").state == str(
+        FILTER_REPLACE_TIME - 743
+    )