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

@@ -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

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

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,

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

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"
+  ]
 }

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
 

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
 

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
 

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
+)

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