Fix roborock image crashes (#116487)

This commit is contained in:
Luke Lashley 2024-05-01 05:56:02 -04:00 committed by GitHub
parent 8230bfcf8f
commit 835ce919f4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 92 additions and 7 deletions

View File

@ -49,7 +49,7 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
) )
device_data = DeviceData(device, product_info.model, device_networking.ip) device_data = DeviceData(device, product_info.model, device_networking.ip)
self.api: RoborockLocalClientV1 | RoborockMqttClientV1 = RoborockLocalClientV1( self.api: RoborockLocalClientV1 | RoborockMqttClientV1 = RoborockLocalClientV1(
device_data device_data, queue_timeout=5
) )
self.cloud_api = cloud_api self.cloud_api = cloud_api
self.device_info = DeviceInfo( self.device_info = DeviceInfo(

View File

@ -66,17 +66,26 @@ class RoborockMap(RoborockCoordinatedEntity, ImageEntity):
) )
self._attr_image_last_updated = dt_util.utcnow() self._attr_image_last_updated = dt_util.utcnow()
self.map_flag = map_flag 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 self._attr_entity_category = EntityCategory.DIAGNOSTIC
@property
def available(self):
"""Determines if the entity is available."""
return self.cached_map != b""
@property @property
def is_selected(self) -> bool: def is_selected(self) -> bool:
"""Return if this map is the currently selected map.""" """Return if this map is the currently selected map."""
return self.map_flag == self.coordinator.current_map return self.map_flag == self.coordinator.current_map
def is_map_valid(self) -> bool: def is_map_valid(self) -> bool:
"""Update this map if it is the current active map, and the vacuum is cleaning.""" """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 ( return self.cached_map == b"" or (
self.is_selected self.is_selected
and self.image_last_updated is not None and self.image_last_updated is not None
and self.coordinator.roborock_device_info.props.status is not None and self.coordinator.roborock_device_info.props.status is not None
@ -96,7 +105,16 @@ class RoborockMap(RoborockCoordinatedEntity, ImageEntity):
async def async_image(self) -> bytes | None: async def async_image(self) -> bytes | None:
"""Update the image if it is not cached.""" """Update the image if it is not cached."""
if self.is_map_valid(): 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) self.cached_map = self._create_image(map_data)
return self.cached_map return self.cached_map
@ -141,9 +159,10 @@ async def create_coordinator_maps(
await asyncio.sleep(MAP_SLEEP) await asyncio.sleep(MAP_SLEEP)
# Get the map data # Get the map data
map_update = await asyncio.gather( 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( entities.append(
RoborockMap( RoborockMap(
f"{slugify(coord.roborock_device_info.device.duid)}_map_{map_info.name}", f"{slugify(coord.roborock_device_info.device.duid)}_map_{map_info.name}",

View File

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

View File

@ -5,7 +5,12 @@ from datetime import timedelta
from http import HTTPStatus from http import HTTPStatus
from unittest.mock import patch 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.core import HomeAssistant
from homeassistant.setup import async_setup_component
from homeassistant.util import dt as dt_util from homeassistant.util import dt as dt_util
from tests.common import MockConfigEntry, async_fire_time_changed 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) async_fire_time_changed(hass, now)
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs") resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
assert not resp.ok 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