Fix roborock image crashes (home-assistant#116487)

homeassistant/components/roborock/coordinator.py

@@ -49,7 +49,7 @@ def __init__(
         )
         device_data = DeviceData(device, product_info.model, device_networking.ip)
         self.api: RoborockLocalClientV1 | RoborockMqttClientV1 = RoborockLocalClientV1(
-            device_data
+            device_data, queue_timeout=5
         )
         self.cloud_api = cloud_api
         self.device_info = DeviceInfo(

homeassistant/components/roborock/image.py

@@ -66,17 +66,26 @@ def __init__(
         )
         self._attr_image_last_updated = dt_util.utcnow()
         self.map_flag = map_flag
-        self.cached_map = self._create_image(starting_map)
+        try:
+            self.cached_map = self._create_image(starting_map)
+        except HomeAssistantError:
+            # If we failed to update the image on init, we set cached_map to empty bytes so that we are unavailable and can try again later.
+            self.cached_map = b""
         self._attr_entity_category = EntityCategory.DIAGNOSTIC
 
+    @property
+    def available(self):
+        """Determines if the entity is available."""
+        return self.cached_map != b""
+
     @property
     def is_selected(self) -> bool:
         """Return if this map is the currently selected map."""
         return self.map_flag == self.coordinator.current_map
 
     def is_map_valid(self) -> bool:
-        """Update this map if it is the current active map, and the vacuum is cleaning."""
-        return (
+        """Update this map if it is the current active map, and the vacuum is cleaning or if it has never been set at all."""
+        return self.cached_map == b"" or (
             self.is_selected
             and self.image_last_updated is not None
             and self.coordinator.roborock_device_info.props.status is not None
@@ -96,7 +105,16 @@ def _handle_coordinator_update(self):
     async def async_image(self) -> bytes | None:
         """Update the image if it is not cached."""
         if self.is_map_valid():
-            map_data: bytes = await self.cloud_api.get_map_v1()
+            response = await asyncio.gather(
+                *(self.cloud_api.get_map_v1(), self.coordinator.get_rooms()),
+                return_exceptions=True,
+            )
+            if not isinstance(response[0], bytes):
+                raise HomeAssistantError(
+                    translation_domain=DOMAIN,
+                    translation_key="map_failure",
+                )
+            map_data = response[0]
             self.cached_map = self._create_image(map_data)
         return self.cached_map
 
@@ -141,9 +159,10 @@ async def create_coordinator_maps(
             await asyncio.sleep(MAP_SLEEP)
         # Get the map data
         map_update = await asyncio.gather(
-            *[coord.cloud_api.get_map_v1(), coord.get_rooms()]
+            *[coord.cloud_api.get_map_v1(), coord.get_rooms()], return_exceptions=True
         )
-        api_data: bytes = map_update[0]
+        # If we fail to get the map -> We should set it to empty byte, still create it, and set it as unavailable.
+        api_data: bytes = map_update[0] if isinstance(map_update[0], bytes) else b""
         entities.append(
             RoborockMap(
                 f"{slugify(coord.roborock_device_info.device.duid)}_map_{map_info.name}",

tests/components/roborock/conftest.py

@@ -91,6 +91,10 @@ def bypass_api_fixture() -> None:
                 RoomMapping(18, "2362041"),
             ],
         ),
+        patch(
+            "homeassistant.components.roborock.coordinator.RoborockMqttClientV1.get_map_v1",
+            return_value=b"123",
+        ),
     ):
         yield
 

tests/components/roborock/test_image.py

@@ -5,7 +5,12 @@
 from http import HTTPStatus
 from unittest.mock import patch
 
+from roborock import RoborockException
+
+from homeassistant.components.roborock import DOMAIN
+from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
 from homeassistant.util import dt as dt_util
 
 from tests.common import MockConfigEntry, async_fire_time_changed
@@ -82,3 +87,60 @@ async def test_floorplan_image_failed_parse(
         async_fire_time_changed(hass, now)
         resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
     assert not resp.ok
+
+
+async def test_fail_parse_on_startup(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    mock_roborock_entry: MockConfigEntry,
+    bypass_api_fixture,
+) -> None:
+    """Test that if we fail parsing on startup, we create the entity but set it as unavailable."""
+    map_data = copy.deepcopy(MAP_DATA)
+    map_data.image = None
+    with patch(
+        "homeassistant.components.roborock.image.RoborockMapDataParser.parse",
+        return_value=map_data,
+    ):
+        await async_setup_component(hass, DOMAIN, {})
+        await hass.async_block_till_done()
+    assert (
+        image_entity := hass.states.get("image.roborock_s7_maxv_upstairs")
+    ) is not None
+    assert image_entity.state == STATE_UNAVAILABLE
+
+
+async def test_fail_updating_image(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    hass_client: ClientSessionGenerator,
+) -> None:
+    """Test that we handle failing getting the image after it has already been setup.."""
+    client = await hass_client()
+    map_data = copy.deepcopy(MAP_DATA)
+    map_data.image = None
+    now = dt_util.utcnow() + timedelta(seconds=91)
+    # Copy the device prop so we don't override it
+    prop = copy.deepcopy(PROP)
+    prop.status.in_cleaning = 1
+    # Update image, but get none for parse image.
+    with (
+        patch(
+            "homeassistant.components.roborock.image.RoborockMapDataParser.parse",
+            return_value=map_data,
+        ),
+        patch(
+            "homeassistant.components.roborock.coordinator.RoborockLocalClientV1.get_prop",
+            return_value=prop,
+        ),
+        patch(
+            "homeassistant.components.roborock.image.dt_util.utcnow", return_value=now
+        ),
+        patch(
+            "homeassistant.components.roborock.coordinator.RoborockMqttClientV1.get_map_v1",
+            side_effect=RoborockException,
+        ),
+    ):
+        async_fire_time_changed(hass, now)
+        resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
+    assert not resp.ok