Add real-time updates for Roborock wet/dry vacuums (#180493)

homeassistant/components/roborock/binary_sensor.py

@@ -278,7 +278,10 @@ class RoborockBinarySensorEntityA01(RoborockCoordinatedEntityA01, BinarySensorEn
 
     @property
     @override
-    def is_on(self) -> bool:
+    def is_on(self) -> bool | None:
         """Return the value reported by the sensor."""
-        value = self.coordinator.data[self.entity_description.data_protocol]
+        if (
+            value := self.coordinator.data.get(self.entity_description.data_protocol)
+        ) is None:
+            return None
         return self.entity_description.value_fn(value)

homeassistant/components/roborock/coordinator.py

@@ -574,6 +574,7 @@ class RoborockWetDryVacUpdateCoordinator(
         """Initialize."""
         super().__init__(hass, config_entry, device)
         self.api = api
+        self._unsub_update = api.add_update_listener(self._handle_update)
         supported_schema_ids = device.product.supported_schema_ids
         self.request_protocols = [
             protocol
@@ -586,13 +587,41 @@ class RoborockWetDryVacUpdateCoordinator(
         self,
     ) -> dict[RoborockDyadDataProtocol, StateType]:
         try:
-            return await self.api.query_values(self.request_protocols)
+            await self.api.query_values(self.request_protocols)
         except RoborockException as ex:
             _LOGGER.debug("Failed to update wet dry vac data: %s", ex)
-            raise UpdateFailed(
-                translation_domain=DOMAIN,
-                translation_key="update_data_fail",
-            ) from ex
+            if not self._should_suppress_update_failure():
+                raise UpdateFailed(
+                    translation_domain=DOMAIN,
+                    translation_key="update_data_fail",
+                ) from ex
+        return self.api.values
+
+    def _should_suppress_update_failure(self) -> bool:
+        """Determine if we should suppress update failure reporting.
+
+        The device leaves the network while it sleeps on its dock, so a poll can
+        fail while the device is still reporting its state on its own.
+        """
+        if (last_message_time := self.api.last_message_time) is None:
+            return False
+        failure_duration = dt_util.utcnow() - last_message_time
+        _LOGGER.debug("Update failure duration: %s", failure_duration)
+        return failure_duration < MIN_UNAVAILABLE_DURATION
+
+    @callback
+    def _handle_update(self) -> None:
+        """Apply the state the device reported on its own."""
+        _LOGGER.debug("Wet dry vac state updated, updating coordinator data")
+        self.data = self.api.values
+        self.last_update_success = True
+        self.async_update_listeners()
+
+    @override
+    async def async_shutdown(self) -> None:
+        """Stop following the device state on shutdown."""
+        self._unsub_update()
+        await super().async_shutdown()
 
 
 class RoborockDataUpdateCoordinatorB01(DataUpdateCoordinator[B01Props]):

homeassistant/components/roborock/sensor.py

@@ -698,7 +698,7 @@ class RoborockSensorEntityA01(RoborockCoordinatedEntityA01, SensorEntity):
     @override
     def native_value(self) -> StateType:
         """Return the value reported by the sensor."""
-        return self.coordinator.data[self.entity_description.data_protocol]
+        return self.coordinator.data.get(self.entity_description.data_protocol)
 
 
 class RoborockSensorEntityB01Q7(RoborockCoordinatedEntityB01Q7, SensorEntity):

tests/components/roborock/conftest.py

@@ -102,17 +102,22 @@ from tests.common import MockConfigEntry
 _LOGGER = logging.getLogger(__name__)
 
 
+DYAD_VALUES: dict[RoborockDyadDataProtocol, Any] = {
+    RoborockDyadDataProtocol.STATUS: RoborockDyadStateCode.drying.name,
+    RoborockDyadDataProtocol.POWER: 100,
+    RoborockDyadDataProtocol.MESH_LEFT: 111,
+    RoborockDyadDataProtocol.BRUSH_LEFT: 222,
+    RoborockDyadDataProtocol.ERROR: DyadError.none.name,
+    RoborockDyadDataProtocol.TOTAL_RUN_TIME: 213,
+}
+
+
 def create_dyad_trait() -> Mock:
     """Create dyad trait for A01 devices."""
     dyad_trait = AsyncMock()
-    dyad_trait.query_values.return_value = {
-        RoborockDyadDataProtocol.STATUS: RoborockDyadStateCode.drying.name,
-        RoborockDyadDataProtocol.POWER: 100,
-        RoborockDyadDataProtocol.MESH_LEFT: 111,
-        RoborockDyadDataProtocol.BRUSH_LEFT: 222,
-        RoborockDyadDataProtocol.ERROR: DyadError.none.name,
-        RoborockDyadDataProtocol.TOTAL_RUN_TIME: 213,
-    }
+    dyad_trait.add_update_listener = Mock(return_value=Mock())
+    dyad_trait.values = dict(DYAD_VALUES)
+    dyad_trait.last_message_time = None
     return dyad_trait
 
 

tests/components/roborock/test_binary_sensor.py

@@ -7,11 +7,12 @@ import pytest
 from roborock.data import RoborockDockTypeCode
 from roborock.device_features import RoborockDockFeatures
 from roborock.exceptions import RoborockException
+from roborock.roborock_message import RoborockZeoProtocol
 from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.components.automation import DOMAIN as AUTOMATION_DOMAIN
 from homeassistant.components.roborock.const import DOMAIN
-from homeassistant.const import STATE_UNAVAILABLE, Platform
+from homeassistant.const import STATE_UNAVAILABLE, STATE_UNKNOWN, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er, issue_registry as ir
 from homeassistant.setup import async_setup_component
@@ -129,6 +130,22 @@ async def test_zeo_request_protocols_filtered_by_schema(
     assert hass.states.get("binary_sensor.zeo_two_softener") is None
 
 
+async def test_zeo_unreported_protocol_is_unknown(
+    hass: HomeAssistant,
+    mock_roborock_entry: MockConfigEntry,
+    fake_devices: list[FakeDevice],
+) -> None:
+    """Test a protocol the device has not reported yet reads as unknown."""
+    zeo = next(device.zeo for device in fake_devices if device.zeo is not None)
+    del zeo.query_values.return_value[RoborockZeoProtocol.DETERGENT_EMPTY]
+
+    await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("binary_sensor.zeo_one_detergent").state == STATE_UNKNOWN
+    assert hass.states.get("binary_sensor.zeo_one_softener").state == "off"
+
+
 @pytest.fixture
 def dock_type(request: pytest.FixtureRequest, fake_vacuum: FakeDevice) -> None:
     """Report the parametrized dock type for the fake vacuum."""

tests/components/roborock/test_sensor.py

@@ -1,21 +1,26 @@
 """Test Roborock Sensors."""
 
+from datetime import timedelta
 from typing import Any
 
+from freezegun.api import FrozenDateTimeFactory
 import pytest
 from roborock.data.v1 import RoborockDockTypeCode
 from roborock.device_features import RoborockDockFeatures
 from roborock.exceptions import RoborockException
+from roborock.roborock_message import RoborockDyadDataProtocol
 from syrupy.assertion import SnapshotAssertion
 
-from homeassistant.components.roborock.const import DOMAIN
-from homeassistant.const import STATE_UNAVAILABLE, Platform
+from homeassistant.components.roborock.const import A01_UPDATE_INTERVAL, DOMAIN
+from homeassistant.components.roborock.coordinator import MIN_UNAVAILABLE_DURATION
+from homeassistant.const import STATE_UNAVAILABLE, STATE_UNKNOWN, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
+from homeassistant.util import dt as dt_util
 
 from .conftest import FakeDevice
 
-from tests.common import MockConfigEntry, snapshot_platform
+from tests.common import MockConfigEntry, async_fire_time_changed, snapshot_platform
 
 
 @pytest.fixture
@@ -156,3 +161,98 @@ async def test_dock_cleaning_brush_sensor_created_when_supported(
     state = hass.states.get("sensor.roborock_s7_maxv_dock_maintenance_brush_time_left")
     assert state is not None
     assert state.state == "235"
+
+
+async def test_dyad_follows_reported_state(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    fake_devices: list[FakeDevice],
+) -> None:
+    """Test the device state is applied as the library reports it."""
+    dyad = next(device.dyad for device in fake_devices if device.dyad is not None)
+    assert hass.states.get("sensor.dyad_pro_battery").state == "100"
+
+    dyad.values = {**dyad.values, RoborockDyadDataProtocol.POWER: 50}
+    dyad.add_update_listener.call_args[0][0]()
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.dyad_pro_battery").state == "50"
+
+
+async def test_dyad_unsubscribed_on_unload(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    fake_devices: list[FakeDevice],
+) -> None:
+    """Test the update listener is removed when the config entry unloads."""
+    dyad = next(device.dyad for device in fake_devices if device.dyad is not None)
+    unsub = dyad.add_update_listener.return_value
+
+    assert await hass.config_entries.async_unload(setup_entry.entry_id)
+    await hass.async_block_till_done()
+
+    unsub.assert_called_once()
+
+
+async def test_dyad_unreported_protocol_is_unknown(
+    hass: HomeAssistant,
+    fake_devices: list[FakeDevice],
+    mock_roborock_entry: MockConfigEntry,
+) -> None:
+    """Test a protocol the device has not reported yet reads as unknown."""
+    dyad = next(device.dyad for device in fake_devices if device.dyad is not None)
+    dyad.values = {RoborockDyadDataProtocol.POWER: 50}
+
+    await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.dyad_pro_battery").state == "50"
+    assert hass.states.get("sensor.dyad_pro_status").state == STATE_UNKNOWN
+
+
+async def test_dyad_update_does_not_postpone_poll(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    fake_devices: list[FakeDevice],
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test the fallback poll keeps its schedule while the device reports state."""
+    dyad = next(device.dyad for device in fake_devices if device.dyad is not None)
+    dyad.query_values.reset_mock()
+
+    freezer.tick(A01_UPDATE_INTERVAL / 2)
+    dyad.add_update_listener.call_args[0][0]()
+    await hass.async_block_till_done()
+
+    freezer.tick(A01_UPDATE_INTERVAL / 2 + timedelta(seconds=1))
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    assert dyad.query_values.call_count == 1
+
+
+@pytest.mark.parametrize(
+    ("last_message_age", "expected_state"),
+    [
+        pytest.param(timedelta(0), "100", id="still_talking"),
+        pytest.param(MIN_UNAVAILABLE_DURATION, STATE_UNAVAILABLE, id="gone_silent"),
+    ],
+)
+async def test_dyad_availability_follows_last_message(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    fake_devices: list[FakeDevice],
+    freezer: FrozenDateTimeFactory,
+    last_message_age: timedelta,
+    expected_state: str,
+) -> None:
+    """Test a failed poll only reports unavailable once the device stops talking."""
+    dyad = next(device.dyad for device in fake_devices if device.dyad is not None)
+    dyad.query_values.side_effect = RoborockException("Simulated failure")
+    dyad.last_message_time = dt_util.utcnow() - last_message_age
+
+    freezer.tick(A01_UPDATE_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.dyad_pro_battery").state == expected_state