Add map image for Roborock B01 Q7 devices (#183884)

homeassistant/components/roborock/coordinator.py

@@ -709,6 +709,7 @@ class RoborockB01Q7UpdateCoordinator(RoborockDataUpdateCoordinatorB01):
                 translation_domain=DOMAIN,
                 translation_key="update_data_fail",
             )
+        await self.async_refresh_q7_map()
         return data
 
     async def async_refresh_q7_map(self) -> bool:

homeassistant/components/roborock/image.py

@@ -1,4 +1,4 @@
-"""Support for Roborock image."""
+"""Define Roborock image entities."""
 
 from datetime import datetime
 import logging
@@ -6,6 +6,7 @@ from typing import override
 
 from roborock.devices.traits.v1.home import HomeTrait
 from roborock.devices.traits.v1.map_content import MapContent
+from roborock.exceptions import RoborockException
 
 from homeassistant.components.image import ImageEntity
 from homeassistant.config_entries import ConfigEntry
@@ -223,43 +224,63 @@ class RoborockMapQ10(RoborockCoordinatedEntityB01Q10, ImageEntity):
 
 
 class RoborockMapQ7(RoborockCoordinatedEntityB01Q7, ImageEntity):
-    """A class to let you visualize the current map of a Q7 device."""
+    """A class to visualize the current map of a Q7 device.
+
+    `python-roborock` here is pinned to >=5.31.1, which predates the
+    unsolicited-map-push support added in 7.1.0 (see README) — so unlike the
+    Q10 equivalent this entity does not rely on push updates arriving on
+    their own. It does an explicit `map.refresh()` + `map_content.refresh()`
+    once when added to hass so an image is available immediately, and again
+    whenever the coordinator's periodic status poll completes while a clean
+    is in progress (map content only meaningfully changes during cleaning).
+    """
 
     _attr_content_type = "image/png"
     _attr_entity_category = EntityCategory.DIAGNOSTIC
     _attr_translation_key = "map"
 
     def __init__(self, coordinator: RoborockB01Q7UpdateCoordinator) -> None:
-        """Initialize the Roborock Q7 map."""
+        """Initialize a Roborock Q7 map image entity."""
         RoborockCoordinatedEntityB01Q7.__init__(
             self, f"map_{coordinator.duid_slug}", coordinator
         )
         ImageEntity.__init__(self, coordinator.hass)
         self._map_content_trait = coordinator.api.map_content
         self._cached_map: bytes | None = None
-        self._attr_image_last_updated = None
 
     @override
     async def async_added_to_hass(self) -> None:
-        """When entity is added to hass, fetch and cache the current map image."""
+        """Register a trait listener and fetch an initial map."""
         await super().async_added_to_hass()
-        await self.coordinator.async_refresh_q7_map()
-        image_content = self._map_content_trait.image_content
-        if image_content is not None:
-            self._cached_map = image_content
-            self._attr_image_last_updated = dt_util.utcnow()
-            self.async_write_ha_state()
+        self.async_on_remove(
+            self._map_content_trait.add_update_listener(self._handle_map_update)
+        )
+        try:
+            await self.coordinator.api.map.refresh()
+            await self._map_content_trait.refresh()
+        except RoborockException as ex:
+            _LOGGER.debug(
+                "Initial Q7 map fetch failed (will retry on next poll): %s", ex
+            )
+        self._handle_map_update()
 
     @callback
     @override
     def _handle_coordinator_update(self) -> None:
-        """Handle updated data from the coordinator."""
-        image_content = self._map_content_trait.image_content
-        if image_content is not None and image_content != self._cached_map:
-            self._cached_map = image_content
-            self._attr_image_last_updated = dt_util.utcnow()
+        """Handle a coordinator update."""
+        self._handle_map_update()
         super()._handle_coordinator_update()
 
+    @callback
+    def _handle_map_update(self) -> None:
+        """Cache the newly fetched/pushed map if its content changed."""
+        image_content = self._map_content_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."""

tests/components/roborock/conftest.py

@@ -181,19 +181,8 @@ def create_b01_q7_trait() -> Mock:
     b01_trait.set_clean_path_preference = AsyncMock()
     b01_trait.set_water_level = AsyncMock()
     b01_trait.send = AsyncMock()
-
-    b01_trait.map = AsyncMock()
-    b01_trait.map.refresh = AsyncMock()
-    b01_trait.map.map_list = []
-    b01_trait.map.current_map_id = None
-
-    b01_trait.map_content = AsyncMock()
-    b01_trait.map_content.refresh = AsyncMock()
-    b01_trait.map_content.image_content = None
-    b01_trait.map_content.map_data = None
-
-    b01_trait.clean_segments = AsyncMock()
-
+    b01_trait.map_content.image_content = b"\x89PNG-q7"
+    attach_update_listeners(b01_trait.map_content)
     return b01_trait
 
 

tests/components/roborock/test_image.py

@@ -13,6 +13,7 @@ from roborock.data import CombinedMapInfo, RoborockStateCode
 from roborock.devices.traits.v1.map_content import MapContent
 
 from homeassistant.components.roborock.const import V1_LOCAL_NOT_CLEANING_INTERVAL
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.util import dt as dt_util
@@ -47,7 +48,7 @@ async def test_floorplan_image(
 ) -> None:
     """Test floor plan map image is correctly set up."""
     assert len(hass.states.async_all("image")) == 6
-
+    assert hass.states.get("image.roborock_q7_map") is not None
     assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None
     # Load the image on demand
     client = await hass_client()
@@ -191,10 +192,10 @@ async def test_map_status_change(
             {
                 "image.roborock_s7_2_downstairs",
                 "image.roborock_s7_2_upstairs",
+                "image.roborock_q7_map",
                 # Expect default names based on map flags
                 "image.roborock_s7_maxv_map_0",
                 "image.roborock_s7_maxv_map_1",
-                "image.roborock_q7_map",
                 "image.roborock_q10_s5_map",
             },
         ),
@@ -278,6 +279,37 @@ async def test_q10_map_image(
     assert await resp.read() == b"\x89PNG-q10-new"
 
 
+async def test_q7_map_image_push_update(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    hass_client: ClientSessionGenerator,
+    fake_devices: list[FakeDevice],
+) -> None:
+    """Test the Q7 map image is updated from a push notification."""
+    entity_id = "image.roborock_q7_map"
+    assert hass.states.get(entity_id) is not None
+
+    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-q7"
+
+    q7_devices = [
+        device for device in fake_devices if device.b01_q7_properties is not None
+    ]
+    assert len(q7_devices) == 1
+    map_content = q7_devices[0].b01_q7_properties.map_content
+    map_content.image_content = b"\x89PNG-q7-new"
+
+    # Simulate a map-content update from the device.
+    map_content.add_update_listener.call_args.args[0]()
+    await hass.async_block_till_done()
+
+    resp = await client.get(f"/api/image_proxy/{entity_id}")
+    assert resp.status == HTTPStatus.OK
+    assert await resp.read() == b"\x89PNG-q7-new"
+
+
 async def test_map_load_delayed(
     hass: HomeAssistant,
     mock_roborock_entry: MockConfigEntry,
@@ -362,14 +394,14 @@ async def test_map_load_delayed(
     assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None
 
 
-async def test_q7_map_image(
+async def test_q7_map_image_coordinator_update(
     hass: HomeAssistant,
     setup_entry: MockConfigEntry,
     hass_client: ClientSessionGenerator,
     fake_q7_vacuum: FakeDevice,
     freezer: FrozenDateTimeFactory,
 ) -> None:
-    """Test the coordinator-driven Q7 map image."""
+    """Test the Q7 map image is updated by the coordinator."""
     entity_id = "image.roborock_q7_map"
     assert hass.states.get(entity_id) is not None
 
@@ -421,3 +453,23 @@ async def test_q7_map_image(
     resp = await client.get(f"/api/image_proxy/{entity_id}")
     assert resp.status == HTTPStatus.OK
     assert await resp.read() == b"\x89PNG-q7-new"
+
+
+async def test_q7_map_initial_fetch_exception(
+    hass: HomeAssistant,
+    mock_roborock_entry: MockConfigEntry,
+    fake_q7_vacuum: FakeDevice,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test Q7 map initial fetch handles RoborockException gracefully."""
+    caplog.set_level(logging.DEBUG, logger="homeassistant.components.roborock")
+    assert fake_q7_vacuum.b01_q7_properties is not None
+
+    fake_q7_vacuum.b01_q7_properties.map.refresh.side_effect = RoborockException("boom")
+
+    await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    assert mock_roborock_entry.state is ConfigEntryState.LOADED
+    assert hass.states.get("image.roborock_q7_map") is not None
+    assert "Initial Q7 map fetch failed (will retry on next poll)" in caplog.text