Update Roborock entities to handle unavailable data (home-assistant#165618)

Author: allenporter

homeassistant/components/roborock/binary_sensor.py

@@ -159,7 +159,11 @@ async def async_setup_entry(
         )
         for coordinator in config_entry.runtime_data.v1
         for description in BINARY_SENSOR_DESCRIPTIONS
-        if description.value_fn(coordinator.data) is not None
+        # Note: Currently coordinator.data is always available on startup but won't be in the future
+        if (
+            coordinator.data is not None
+            and description.value_fn(coordinator.data) is not None
+        )
     ]
     entities.extend(
         RoborockBinarySensorEntityA01(
@@ -193,9 +197,11 @@ def __init__(
         self.entity_description = description
 
     @property
-    def is_on(self) -> bool:
+    def is_on(self) -> bool | None:
         """Return the value reported by the sensor."""
-        return bool(self.entity_description.value_fn(self.coordinator.data))
+        if (data := self.coordinator.data) is not None:
+            return bool(self.entity_description.value_fn(data))
+        return None
 
 
 class RoborockBinarySensorEntityA01(RoborockCoordinatedEntityA01, BinarySensorEntity):

homeassistant/components/roborock/coordinator.py

@@ -83,7 +83,7 @@ def values(
 type RoborockConfigEntry = ConfigEntry[RoborockCoordinators]
 
 
-class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceState]):
+class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceState | None]):
     """Class to manage fetching data from the API."""
 
     config_entry: RoborockConfigEntry
@@ -229,7 +229,7 @@ async def _update_device_prop(self) -> None:
         )
         _LOGGER.debug("Updated device properties")
 
-    async def _async_update_data(self) -> DeviceState:
+    async def _async_update_data(self) -> DeviceState | None:
         """Update data via library."""
         await self._verify_api()
         try:

homeassistant/components/roborock/entity.py

@@ -3,7 +3,6 @@
 from typing import Any
 
 from roborock.devices.traits.v1.command import CommandTrait
-from roborock.devices.traits.v1.status import StatusTrait
 from roborock.exceptions import RoborockException
 from roborock.roborock_typing import RoborockCommand
 
@@ -94,12 +93,6 @@ def __init__(
         CoordinatorEntity.__init__(self, coordinator=coordinator)
         self._attr_unique_id = unique_id
 
-    @property
-    def _device_status(self) -> StatusTrait:
-        """Return the status of the device."""
-        data = self.coordinator.data
-        return data.status
-
     async def send(
         self,
         command: RoborockCommand | str,

homeassistant/components/roborock/image.py

@@ -9,7 +9,7 @@
 from homeassistant.components.image import ImageEntity
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
@@ -73,6 +73,7 @@ def __init__(
         self.map_flag = map_flag
         self.cached_map: bytes | None = None
         self._attr_entity_category = EntityCategory.DIAGNOSTIC
+        self._attr_image_last_updated = None
 
     async def async_added_to_hass(self) -> None:
         """When entity is added to hass load any previously cached maps from disk."""
@@ -88,17 +89,17 @@ def _map_content(self) -> MapContent | None:
             return map_content
         return None
 
+    @callback
     def _handle_coordinator_update(self) -> None:
         """Handle updated data from the coordinator.
 
         If the coordinator has updated the map, we can update the image.
         """
-        if (map_content := self._map_content) is None:
+        if self.coordinator.data is None or (map_content := self._map_content) is None:
             return
         if self.cached_map != map_content.image_content:
             self.cached_map = map_content.image_content
             self._attr_image_last_updated = self.coordinator.last_home_update
-
         super()._handle_coordinator_update()
 
     async def async_image(self) -> bytes | None:

homeassistant/components/roborock/sensor.py

@@ -427,7 +427,11 @@ async def async_setup_entry(
         )
         for coordinator in coordinators.v1
         for description in SENSOR_DESCRIPTIONS
-        if description.value_fn(coordinator.data) is not None
+        # Note: Currently coordinator.data is always available on startup but won't be in the future
+        if (
+            coordinator.data is not None
+            and description.value_fn(coordinator.data) is not None
+        )
     ]
     entities.extend(RoborockCurrentRoom(coordinator) for coordinator in coordinators.v1)
     entities.extend(
@@ -480,6 +484,8 @@ def __init__(
     @property
     def native_value(self) -> StateType | datetime.datetime:
         """Return the value reported by the sensor."""
+        if self.coordinator.data is None:
+            return None
         return self.entity_description.value_fn(self.coordinator.data)
 
 

homeassistant/components/roborock/vacuum.py

@@ -150,6 +150,7 @@ def __init__(
             coordinator.duid_slug,
             coordinator,
         )
+        self._status_trait = coordinator.properties_api.status
         self._home_trait = coordinator.properties_api.home
         self._maps_trait = coordinator.properties_api.maps
 
@@ -176,31 +177,37 @@ def _handle_coordinator_update(self) -> None:
     @property
     def fan_speed_list(self) -> list[str]:
         """Get the list of available fan speeds."""
-        return [mode.value for mode in self._device_status.fan_speed_options]
+        if self.coordinator.data is None:
+            return []
+        return [mode.value for mode in self._status_trait.fan_speed_options]
 
     @property
     def activity(self) -> VacuumActivity | None:
         """Return the status of the vacuum cleaner."""
-        assert self._device_status.state is not None
-        return STATE_CODE_TO_STATE.get(self._device_status.state)
+        if self.coordinator.data is None or self._status_trait.state is None:
+            return None
+        return STATE_CODE_TO_STATE.get(self._status_trait.state)
 
     @property
     def fan_speed(self) -> str | None:
         """Return the fan speed of the vacuum cleaner."""
-        return self._device_status.fan_speed_name
+        if self.coordinator.data is None:
+            return None
+        return self._status_trait.fan_speed_name
 
     async def async_start(self) -> None:
         """Start the vacuum."""
-        if self._device_status.in_returning == 1:
-            await self.send(RoborockCommand.APP_CHARGE)
-        elif self._device_status.in_cleaning == 2:
-            await self.send(RoborockCommand.RESUME_ZONED_CLEAN)
-        elif self._device_status.in_cleaning == 3:
-            await self.send(RoborockCommand.RESUME_SEGMENT_CLEAN)
-        elif self._device_status.in_cleaning == 4:
-            await self.send(RoborockCommand.APP_RESUME_BUILD_MAP)
-        else:
-            await self.send(RoborockCommand.APP_START)
+        command = RoborockCommand.APP_START
+        if self.coordinator.data is not None:
+            if self._status_trait.in_returning == 1:
+                command = RoborockCommand.APP_CHARGE
+            elif self._status_trait.in_cleaning == 2:
+                command = RoborockCommand.RESUME_ZONED_CLEAN
+            elif self._status_trait.in_cleaning == 3:
+                command = RoborockCommand.RESUME_SEGMENT_CLEAN
+            elif self._status_trait.in_cleaning == 4:
+                command = RoborockCommand.APP_RESUME_BUILD_MAP
+        await self.send(command)
 
     async def async_pause(self) -> None:
         """Pause the vacuum."""
@@ -224,10 +231,15 @@ async def async_locate(self, **kwargs: Any) -> None:
 
     async def async_set_fan_speed(self, fan_speed: str, **kwargs: Any) -> None:
         """Set vacuum fan speed."""
+        if self.coordinator.data is None:
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="update_options_failed",
+            )
         await self.send(
             RoborockCommand.SET_CUSTOM_MODE,
             [
-                {v: k for k, v in self._device_status.fan_speed_mapping.items()}[
+                {v: k for k, v in self._status_trait.fan_speed_mapping.items()}[
                     fan_speed
                 ]
             ],