Add Roborock Q10 map image entity (#173883)

homeassistant/components/roborock/icons.json

@@ -40,6 +40,11 @@
         "default": "mdi:brush"
       }
     },
+    "image": {
+      "map": {
+        "default": "mdi:floor-plan"
+      }
+    },
     "number": {
       "volume": {
         "default": "mdi:volume-source"

homeassistant/components/roborock/image.py

@@ -14,13 +14,15 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.util import dt as dt_util
 
 from .coordinator import (
+    RoborockB01Q10UpdateCoordinator,
     RoborockConfigEntry,
     RoborockCoordinatorType,
     RoborockDataUpdateCoordinator,
 )
-from .entity import RoborockCoordinatedEntityV1
+from .entity import RoborockCoordinatedEntityB01Q10, RoborockCoordinatedEntityV1
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -40,20 +42,22 @@ async def async_setup_entry(
         coordinator: RoborockCoordinatorType,
     ) -> None:
         """Add entities for a specific coordinator."""
-        if not isinstance(coordinator, RoborockDataUpdateCoordinator):
-            return
-        entities = [
-            RoborockMap(
-                config_entry,
-                coordinator,
-                coordinator.properties_api.home,
-                map_info.map_flag,
-                map_info.name,
+        entities: list[ImageEntity] = []
+        if isinstance(coordinator, RoborockDataUpdateCoordinator):
+            entities.extend(
+                RoborockMap(
+                    config_entry,
+                    coordinator,
+                    coordinator.properties_api.home,
+                    map_info.map_flag,
+                    map_info.name,
+                )
+                for map_info in (
+                    coordinator.properties_api.home.home_map_info or {}
+                ).values()
             )
-            for map_info in (
-                coordinator.properties_api.home.home_map_info or {}
-            ).values()
-        ]
+        elif isinstance(coordinator, RoborockB01Q10UpdateCoordinator):
+            entities.append(RoborockMapQ10(coordinator))
         async_add_entities(entities)
 
     for coordinator in coordinators.values():
@@ -134,3 +138,50 @@ class RoborockMap(RoborockCoordinatedEntityV1, ImageEntity):
         if (map_content := self._map_content) is None:
             raise HomeAssistantError("Map flag not found in coordinator maps")
         return map_content.image_content
+
+
+class RoborockMapQ10(RoborockCoordinatedEntityB01Q10, ImageEntity):
+    """A class to let you visualize the current map of a Q10 device.
+
+    The Q10 pushes its current map over MQTT rather than serving it on
+    request, and the multi-map list is not reachable on this channel, so the
+    device exposes a single push-driven map entity.
+    """
+
+    _attr_content_type = "image/png"
+    _attr_entity_category = EntityCategory.DIAGNOSTIC
+    _attr_translation_key = "map"
+
+    def __init__(self, coordinator: RoborockB01Q10UpdateCoordinator) -> None:
+        """Initialize a Roborock Q10 map."""
+        RoborockCoordinatedEntityB01Q10.__init__(
+            self, f"map_{coordinator.duid_slug}", coordinator
+        )
+        ImageEntity.__init__(self, coordinator.hass)
+        self._map_trait = coordinator.api.map
+        self._cached_map: bytes | None = None
+
+    @override
+    async def async_added_to_hass(self) -> None:
+        """Register a trait listener for push-based map updates."""
+        await super().async_added_to_hass()
+        self.async_on_remove(
+            self._map_trait.add_update_listener(self._handle_map_update)
+        )
+        # Pick up a map that was pushed before the entity was added.
+        self._handle_map_update()
+
+    @callback
+    def _handle_map_update(self) -> None:
+        """Cache the newly pushed map if its content changed."""
+        image_content = self._map_trait.image_content
+        if image_content is None or image_content == self._cached_map:
+            return
+        self._cached_map = image_content
+        self._attr_image_last_updated = dt_util.utcnow()
+        self.async_write_ha_state()
+
+    @override
+    async def async_image(self) -> bytes | None:
+        """Get the cached image."""
+        return self._cached_map

homeassistant/components/roborock/strings.json

@@ -125,6 +125,11 @@
         "name": "Start"
       }
     },
+    "image": {
+      "map": {
+        "name": "Map"
+      }
+    },
     "number": {
       "volume": {
         "name": "Volume"

tests/components/roborock/conftest.py

@@ -239,6 +239,7 @@ def create_b01_q10_trait() -> Mock:
     q10_trait.button_light.disable = AsyncMock()
 
     q10_trait.map = Mock()
+    q10_trait.map.image_content = b"\x89PNG-q10"
     q10_trait.map.rooms = [
         Q10Room(id=9, raw_name="rr_bedroom", pixel_value=36, pixel_count=100),
         Q10Room(id=10, raw_name="rr_living_room", pixel_value=40, pixel_count=200),

tests/components/roborock/test_image.py

@@ -6,6 +6,7 @@ from http import HTTPStatus
 import logging
 from unittest.mock import patch
 
+from freezegun.api import FrozenDateTimeFactory
 import pytest
 from roborock import MultiMapsList, RoborockException
 from roborock.data import RoborockStateCode
@@ -45,7 +46,7 @@ async def test_floorplan_image(
     fake_devices: list[FakeDevice],
 ) -> None:
     """Test floor plan map image is correctly set up."""
-    assert len(hass.states.async_all("image")) == 4
+    assert len(hass.states.async_all("image")) == 5
 
     assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None
     # Load the image on demand
@@ -131,7 +132,7 @@ async def test_map_status_change(
     fake_vacuum: FakeDevice,
 ) -> None:
     """Test floor plan map image is correctly updated on status change."""
-    assert len(hass.states.async_all("image")) == 4
+    assert len(hass.states.async_all("image")) == 5
 
     assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None
     client = await hass_client()
@@ -181,6 +182,7 @@ async def test_map_status_change(
                 "image.roborock_s7_2_upstairs",
                 "image.roborock_s7_maxv_downstairs",
                 "image.roborock_s7_maxv_upstairs",
+                "image.roborock_q10_s5_map",
             },
         ),
         (
@@ -191,6 +193,7 @@ async def test_map_status_change(
                 # Expect default names based on map flags
                 "image.roborock_s7_maxv_map_0",
                 "image.roborock_s7_maxv_map_1",
+                "image.roborock_q10_s5_map",
             },
         ),
     ],
@@ -222,3 +225,52 @@ async def test_image_entity_naming(
     assert {
         state.entity_id for state in hass.states.async_all("image")
     } == expected_entity_ids
+
+
+async def test_q10_map_image(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    hass_client: ClientSessionGenerator,
+    fake_q10_vacuum: FakeDevice,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test the push-driven Q10 map image."""
+    entity_id = "image.roborock_q10_s5_map"
+    assert hass.states.get(entity_id) is not None
+
+    # The map pushed before startup is served
+    client = await hass_client()
+    resp = await client.get(f"/api/image_proxy/{entity_id}")
+    assert resp.status == HTTPStatus.OK
+    assert await resp.read() == b"\x89PNG-q10"
+
+    assert fake_q10_vacuum.b01_q10_properties is not None
+    map_trait = fake_q10_vacuum.b01_q10_properties.map
+
+    def push_update() -> None:
+        for call in map_trait.add_update_listener.call_args_list:
+            call.args[0]()
+
+    # A push that does not change the map content must not update the entity
+    state = hass.states.get(entity_id)
+    assert state is not None
+    last_updated = state.state
+    freezer.tick(timedelta(seconds=30))
+    push_update()
+    await hass.async_block_till_done()
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == last_updated
+
+    # The device pushes an updated map
+    freezer.tick(timedelta(seconds=30))
+    map_trait.image_content = b"\x89PNG-q10-new"
+    push_update()
+    await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state != last_updated
+    resp = await client.get(f"/api/image_proxy/{entity_id}")
+    assert resp.status == HTTPStatus.OK
+    assert await resp.read() == b"\x89PNG-q10-new"