From bee457ed6f6dd623fd785169a0815f04deb3f192 Mon Sep 17 00:00:00 2001 From: Luke Lashley Date: Sat, 18 Nov 2023 15:22:30 -0500 Subject: [PATCH] Add Image to Roborock to display maps (#102941) * add image to roborock * add vacuum position * addressing MR comments * remove room names as it isn't supported in base package * 100% coverage * remove unneeded map changes * fix image logic * optimize create_coordinator_maps * only update time if map is valid * Update test_image.py * fix linting from merge conflict * fix mypy complaints * re-add vacuum to const * fix hanging test * Make map sleep a const * adjust commenting to be less than 88 characters. * bump map parser --- homeassistant/components/roborock/const.py | 13 ++ .../components/roborock/coordinator.py | 13 ++ homeassistant/components/roborock/device.py | 6 + homeassistant/components/roborock/image.py | 151 ++++++++++++++++++ .../components/roborock/manifest.json | 5 +- requirements_all.txt | 3 + requirements_test_all.txt | 3 + tests/components/roborock/conftest.py | 19 ++- tests/components/roborock/mock_data.py | 34 ++++ tests/components/roborock/test_image.py | 75 +++++++++ 10 files changed, 320 insertions(+), 2 deletions(-) create mode 100644 homeassistant/components/roborock/image.py create mode 100644 tests/components/roborock/test_image.py diff --git a/homeassistant/components/roborock/const.py b/homeassistant/components/roborock/const.py index d135f323e90..d7a3a9229f5 100644 --- a/homeassistant/components/roborock/const.py +++ b/homeassistant/components/roborock/const.py @@ -1,4 +1,6 @@ """Constants for Roborock.""" +from vacuum_map_parser_base.config.drawable import Drawable + from homeassistant.const import Platform DOMAIN = "roborock" @@ -9,6 +11,7 @@ CONF_USER_DATA = "user_data" PLATFORMS = [ Platform.BUTTON, Platform.BINARY_SENSOR, + Platform.IMAGE, Platform.NUMBER, Platform.SELECT, Platform.SENSOR, @@ -16,3 +19,13 @@ PLATFORMS = [ Platform.TIME, Platform.VACUUM, ] + +IMAGE_DRAWABLES: list[Drawable] = [ + Drawable.PATH, + Drawable.CHARGER, + Drawable.VACUUM_POSITION, +] + +IMAGE_CACHE_INTERVAL = 90 + +MAP_SLEEP = 3 diff --git a/homeassistant/components/roborock/coordinator.py b/homeassistant/components/roborock/coordinator.py index 30bfc71ea48..cd08cf871d4 100644 --- a/homeassistant/components/roborock/coordinator.py +++ b/homeassistant/components/roborock/coordinator.py @@ -55,6 +55,7 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]): model=self.roborock_device_info.product.model, sw_version=self.roborock_device_info.device.fv, ) + self.current_map: int | None = None if mac := self.roborock_device_info.network_info.mac: self.device_info[ATTR_CONNECTIONS] = {(dr.CONNECTION_NETWORK_MAC, mac)} @@ -91,6 +92,18 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]): """Update data via library.""" try: await self._update_device_prop() + self._set_current_map() except RoborockException as ex: raise UpdateFailed(ex) from ex return self.roborock_device_info.props + + def _set_current_map(self) -> None: + if ( + self.roborock_device_info.props.status is not None + and self.roborock_device_info.props.status.map_status is not None + ): + # The map status represents the map flag as flag * 4 + 3 - + # so we have to invert that in order to get the map flag that we can use to set the current map. + self.current_map = ( + self.roborock_device_info.props.status.map_status - 3 + ) // 4 diff --git a/homeassistant/components/roborock/device.py b/homeassistant/components/roborock/device.py index c8f45b40d82..5fca40a9fd8 100644 --- a/homeassistant/components/roborock/device.py +++ b/homeassistant/components/roborock/device.py @@ -3,6 +3,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.exceptions import RoborockException @@ -82,6 +83,11 @@ class RoborockCoordinatedEntity( data = self.coordinator.data return data.status + @property + def cloud_api(self) -> RoborockMqttClient: + """Return the cloud api.""" + return self.coordinator.cloud_api + async def send( self, command: RoborockCommand | str, diff --git a/homeassistant/components/roborock/image.py b/homeassistant/components/roborock/image.py new file mode 100644 index 00000000000..6c4c7553c14 --- /dev/null +++ b/homeassistant/components/roborock/image.py @@ -0,0 +1,151 @@ +"""Support for Roborock image.""" +import asyncio +import io +from itertools import chain + +from roborock import RoborockCommand +from vacuum_map_parser_base.config.color import ColorsPalette +from vacuum_map_parser_base.config.image_config import ImageConfig +from vacuum_map_parser_base.config.size import Sizes +from vacuum_map_parser_roborock.map_data_parser import RoborockMapDataParser + +from homeassistant.components.image import ImageEntity +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import HomeAssistant +from homeassistant.exceptions import HomeAssistantError +from homeassistant.helpers.entity_platform import AddEntitiesCallback +from homeassistant.util import slugify +import homeassistant.util.dt as dt_util + +from .const import DOMAIN, IMAGE_CACHE_INTERVAL, IMAGE_DRAWABLES, MAP_SLEEP +from .coordinator import RoborockDataUpdateCoordinator +from .device import RoborockCoordinatedEntity + + +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +) -> None: + """Set up Roborock image platform.""" + + coordinators: dict[str, RoborockDataUpdateCoordinator] = hass.data[DOMAIN][ + config_entry.entry_id + ] + entities = list( + chain.from_iterable( + await asyncio.gather( + *(create_coordinator_maps(coord) for coord in coordinators.values()) + ) + ) + ) + async_add_entities(entities) + + +class RoborockMap(RoborockCoordinatedEntity, ImageEntity): + """A class to let you visualize the map.""" + + _attr_has_entity_name = True + + def __init__( + self, + unique_id: str, + coordinator: RoborockDataUpdateCoordinator, + map_flag: int, + starting_map: bytes, + map_name: str, + ) -> None: + """Initialize a Roborock map.""" + RoborockCoordinatedEntity.__init__(self, unique_id, coordinator) + ImageEntity.__init__(self, coordinator.hass) + self._attr_name = map_name + self.parser = RoborockMapDataParser( + ColorsPalette(), Sizes(), IMAGE_DRAWABLES, ImageConfig(), [] + ) + self._attr_image_last_updated = dt_util.utcnow() + self.map_flag = map_flag + self.cached_map = self._create_image(starting_map) + + def is_map_valid(self) -> bool: + """Update this map if it is the current active map, and the vacuum is cleaning.""" + return ( + self.map_flag == self.coordinator.current_map + and self.image_last_updated is not None + and self.coordinator.roborock_device_info.props.status is not None + and bool(self.coordinator.roborock_device_info.props.status.in_cleaning) + ) + + def _handle_coordinator_update(self): + # Bump last updated every third time the coordinator runs, so that async_image + # will be called and we will evaluate on the new coordinator data if we should + # update the cache. + if ( + dt_util.utcnow() - self.image_last_updated + ).total_seconds() > IMAGE_CACHE_INTERVAL and self.is_map_valid(): + self._attr_image_last_updated = dt_util.utcnow() + super()._handle_coordinator_update() + + 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() + self.cached_map = self._create_image(map_data) + return self.cached_map + + def _create_image(self, map_bytes: bytes) -> bytes: + """Create an image using the map parser.""" + parsed_map = self.parser.parse(map_bytes) + if parsed_map.image is None: + raise HomeAssistantError("Something went wrong creating the map.") + img_byte_arr = io.BytesIO() + parsed_map.image.data.save(img_byte_arr, format="PNG") + return img_byte_arr.getvalue() + + +async def create_coordinator_maps( + coord: RoborockDataUpdateCoordinator, +) -> list[RoborockMap]: + """Get the starting map information for all maps for this device. The following steps must be done synchronously. + + Only one map can be loaded at a time per device. + """ + entities = [] + maps = await coord.cloud_api.get_multi_maps_list() + if maps is not None and maps.map_info is not None: + cur_map = coord.current_map + # This won't be None at this point as the coordinator will have run first. + assert cur_map is not None + # Sort the maps so that we start with the current map and we can skip the + # load_multi_map call. + maps_info = sorted( + maps.map_info, key=lambda data: data.mapFlag == cur_map, reverse=True + ) + for roborock_map in maps_info: + # Load the map - so we can access it with get_map_v1 + if roborock_map.mapFlag != cur_map: + # Only change the map and sleep if we have multiple maps. + await coord.api.send_command( + RoborockCommand.LOAD_MULTI_MAP, [roborock_map.mapFlag] + ) + # We cannot get the map until the roborock servers fully process the + # map change. + await asyncio.sleep(MAP_SLEEP) + # Get the map data + api_data: bytes = await coord.cloud_api.get_map_v1() + entities.append( + RoborockMap( + f"{slugify(coord.roborock_device_info.device.duid)}_map_{roborock_map.name}", + coord, + roborock_map.mapFlag, + api_data, + roborock_map.name, + ) + ) + if len(maps.map_info) != 1: + # Set the map back to the map the user previously had selected so that it + # does not change the end user's app. + # Only needs to happen when we changed maps above. + await coord.cloud_api.send_command( + RoborockCommand.LOAD_MULTI_MAP, [cur_map] + ) + return entities diff --git a/homeassistant/components/roborock/manifest.json b/homeassistant/components/roborock/manifest.json index ed043582a0e..3b741995cd4 100644 --- a/homeassistant/components/roborock/manifest.json +++ b/homeassistant/components/roborock/manifest.json @@ -6,5 +6,8 @@ "documentation": "https://www.home-assistant.io/integrations/roborock", "iot_class": "local_polling", "loggers": ["roborock"], - "requirements": ["python-roborock==0.36.1"] + "requirements": [ + "python-roborock==0.36.1", + "vacuum-map-parser-roborock==0.1.1" + ] } diff --git a/requirements_all.txt b/requirements_all.txt index a8ccb26797f..2efdb98189f 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -2671,6 +2671,9 @@ url-normalize==1.4.3 # homeassistant.components.uvc uvcclient==0.11.0 +# homeassistant.components.roborock +vacuum-map-parser-roborock==0.1.1 + # homeassistant.components.vallox vallox-websocket-api==4.0.2 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 89a5db104ae..8b1319ad920 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -1984,6 +1984,9 @@ url-normalize==1.4.3 # homeassistant.components.uvc uvcclient==0.11.0 +# homeassistant.components.roborock +vacuum-map-parser-roborock==0.1.1 + # homeassistant.components.vallox vallox-websocket-api==4.0.2 diff --git a/tests/components/roborock/conftest.py b/tests/components/roborock/conftest.py index 3435bd58cb3..b0a01137ab9 100644 --- a/tests/components/roborock/conftest.py +++ b/tests/components/roborock/conftest.py @@ -12,7 +12,16 @@ from homeassistant.const import CONF_USERNAME from homeassistant.core import HomeAssistant from homeassistant.setup import async_setup_component -from .mock_data import BASE_URL, HOME_DATA, NETWORK_INFO, PROP, USER_DATA, USER_EMAIL +from .mock_data import ( + BASE_URL, + HOME_DATA, + MAP_DATA, + MULTI_MAP_LIST, + NETWORK_INFO, + PROP, + USER_DATA, + USER_EMAIL, +) from tests.common import MockConfigEntry @@ -33,6 +42,12 @@ def bypass_api_fixture() -> None: ), patch( "homeassistant.components.roborock.coordinator.RoborockLocalClient.get_prop", return_value=PROP, + ), patch( + "homeassistant.components.roborock.coordinator.RoborockMqttClient.get_multi_maps_list", + return_value=MULTI_MAP_LIST, + ), patch( + "homeassistant.components.roborock.image.RoborockMapDataParser.parse", + return_value=MAP_DATA, ), patch( "homeassistant.components.roborock.coordinator.RoborockLocalClient.send_message" ), patch( @@ -43,6 +58,8 @@ def bypass_api_fixture() -> None: "roborock.api.AttributeCache.async_value" ), patch( "roborock.api.AttributeCache.value" + ), patch( + "homeassistant.components.roborock.image.MAP_SLEEP", 0 ): yield diff --git a/tests/components/roborock/mock_data.py b/tests/components/roborock/mock_data.py index 87ed02bc3ec..8935a77f142 100644 --- a/tests/components/roborock/mock_data.py +++ b/tests/components/roborock/mock_data.py @@ -1,17 +1,22 @@ """Mock data for Roborock tests.""" from __future__ import annotations +from PIL import Image from roborock.containers import ( CleanRecord, CleanSummary, Consumable, DnDTimer, HomeData, + MultiMapsList, NetworkInfo, S7Status, UserData, ) from roborock.roborock_typing import DeviceProp +from vacuum_map_parser_base.config.image_config import ImageConfig +from vacuum_map_parser_base.map_data import ImageData +from vacuum_map_parser_roborock.map_data_parser import MapData from homeassistant.components.roborock import CONF_BASE_URL, CONF_USER_DATA from homeassistant.const import CONF_USERNAME @@ -418,3 +423,32 @@ PROP = DeviceProp( NETWORK_INFO = NetworkInfo( ip="123.232.12.1", ssid="wifi", mac="ac:cc:cc:cc:cc", bssid="bssid", rssi=90 ) + +MULTI_MAP_LIST = MultiMapsList.from_dict( + { + "maxMultiMap": 4, + "maxBakMap": 1, + "multiMapCount": 2, + "mapInfo": [ + { + "mapFlag": 0, + "addTime": 1686235489, + "length": 8, + "name": "Upstairs", + "bakMaps": [{"addTime": 1673304288}], + }, + { + "mapFlag": 1, + "addTime": 1697579901, + "length": 10, + "name": "Downstairs", + "bakMaps": [{"addTime": 1695521431}], + }, + ], + } +) + +MAP_DATA = MapData(0, 0) +MAP_DATA.image = ImageData( + 100, 10, 10, 10, 10, ImageConfig(), Image.new("RGB", (1, 1)), lambda p: p +) diff --git a/tests/components/roborock/test_image.py b/tests/components/roborock/test_image.py new file mode 100644 index 00000000000..80d4bd37337 --- /dev/null +++ b/tests/components/roborock/test_image.py @@ -0,0 +1,75 @@ +"""Test Roborock Image platform.""" +import copy +from datetime import timedelta +from http import HTTPStatus +from unittest.mock import patch + +from homeassistant.core import HomeAssistant +from homeassistant.util import dt as dt_util + +from tests.common import MockConfigEntry, async_fire_time_changed +from tests.components.roborock.mock_data import MAP_DATA, PROP +from tests.typing import ClientSessionGenerator + + +async def test_floorplan_image( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + hass_client: ClientSessionGenerator, +) -> None: + """Test floor plan map image is correctly set up.""" + # Setup calls the image parsing the first time and caches it. + assert len(hass.states.async_all("image")) == 4 + + assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None + # call a second time -should return cached data + client = await hass_client() + resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs") + assert resp.status == HTTPStatus.OK + body = await resp.read() + assert body is not None + # Call a third time - this time forcing it to update + now = dt_util.utcnow() + timedelta(seconds=91) + async_fire_time_changed(hass, now) + # Copy the device prop so we don't override it + prop = copy.deepcopy(PROP) + prop.status.in_cleaning = 1 + with patch( + "homeassistant.components.roborock.coordinator.RoborockLocalClient.get_prop", + return_value=prop, + ), patch( + "homeassistant.components.roborock.image.dt_util.utcnow", return_value=now + ): + await hass.async_block_till_done() + resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs") + assert resp.status == HTTPStatus.OK + body = await resp.read() + assert body is not None + + +async def test_floorplan_image_failed_parse( + hass: HomeAssistant, + setup_entry: MockConfigEntry, + hass_client: ClientSessionGenerator, +) -> None: + """Test that we correctly handle getting None from the image parser.""" + client = await hass_client() + map_data = copy.deepcopy(MAP_DATA) + map_data.image = None + now = dt_util.utcnow() + timedelta(seconds=91) + async_fire_time_changed(hass, now) + # 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.RoborockLocalClient.get_prop", + return_value=prop, + ), patch( + "homeassistant.components.roborock.image.dt_util.utcnow", return_value=now + ): + resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs") + assert not resp.ok