feat: implement dynamic map update listeners and filtering on HomeTrait

roborock/devices/traits/v1/home.py

@@ -18,10 +18,12 @@
 import asyncio
 import base64
 import logging
+from collections.abc import Callable
 
 from roborock.data import CombinedMapInfo, MultiMapsListMapInfo, NamedRoomMapping, RoborockBase
 from roborock.data.v1.v1_code_mappings import RoborockStateCode
 from roborock.devices.cache import DeviceCache
+from roborock.devices.traits.common import TraitUpdateListener
 from roborock.devices.traits.v1 import common
 from roborock.exceptions import RoborockDeviceBusy, RoborockException, RoborockInvalidStatus
 from roborock.roborock_typing import RoborockCommand
@@ -36,7 +38,7 @@
 MAP_SLEEP = 3
 
 
-class HomeTrait(RoborockBase, common.V1TraitMixin):
+class HomeTrait(RoborockBase, common.V1TraitMixin, TraitUpdateListener):
     """Trait that represents a full view of the home layout."""
 
     command = RoborockCommand.GET_MAP_V1  # This is not used
@@ -66,6 +68,7 @@ def __init__(
         accuracy.
         """
         super().__init__()
+        TraitUpdateListener.__init__(self, logger=_LOGGER)
         self._status_trait = status_trait
         self._maps_trait = maps_trait
         self._map_content = map_content
@@ -100,6 +103,7 @@ async def discover_home(self) -> None:
                 _LOGGER.warning("Failed to parse cached home map content, will re-discover: %s", ex)
                 self._home_map_content = {}
             else:
+                self._notify_update()
                 return
 
         if self._status_trait.state == RoborockStateCode.cleaning:
@@ -202,6 +206,7 @@ async def refresh(self) -> None:
             map_flag := self._maps_trait.current_map
         ) is None:
             _LOGGER.debug("Cannot refresh home data without current map info")
+            self._notify_update()
             return
 
         # Refresh the map content to ensure we have the latest image and object positions
@@ -212,10 +217,26 @@ async def refresh(self) -> None:
             map_flag, combined_map_info, new_map_content, update_cache=self._discovery_completed
         )
 
+    def add_update_listener(self, callback: Callable[[], None]) -> Callable[[], None]:
+        """Register a callback when the trait has been updated.
+
+        Overridden to immediately execute the callback with the current state if populated.
+        """
+        unsub = super().add_update_listener(callback)
+        if self._home_map_info is not None:
+            callback()
+        return unsub
+
     @property
     def home_map_info(self) -> dict[int, CombinedMapInfo] | None:
         """Returns the map information for all cached maps."""
-        return self._home_map_info
+        if self._home_map_info is None or self._maps_trait.map_info is None:
+            return self._home_map_info
+        return {
+            mi.map_flag: value
+            for mi in self._maps_trait.map_info
+            if (value := self._home_map_info.get(mi.map_flag)) is not None
+        }
 
     @property
     def current_map_data(self) -> CombinedMapInfo | None:
@@ -235,7 +256,13 @@ def current_rooms(self) -> list[NamedRoomMapping]:
     @property
     def home_map_content(self) -> dict[int, MapContent] | None:
         """Returns the map content for all cached maps."""
-        return self._home_map_content
+        if self._home_map_content is None or self._maps_trait.map_info is None:
+            return self._home_map_content
+        return {
+            mi.map_flag: value
+            for mi in self._maps_trait.map_info
+            if (value := self._home_map_content.get(mi.map_flag)) is not None
+        }
 
     async def _update_home_cache(
         self, home_map_info: dict[int, CombinedMapInfo], home_map_content: dict[int, MapContent]
@@ -251,6 +278,7 @@ async def _update_home_cache(
         await self._device_cache.set(device_cache_data)
         self._home_map_info = home_map_info
         self._home_map_content = home_map_content
+        self._notify_update()
 
     async def _update_current_map(
         self,
@@ -283,3 +311,4 @@ async def _update_current_map(
         if self._home_map_content is None:
             self._home_map_content = {}
         self._home_map_content[map_flag] = map_content
+        self._notify_update()

tests/devices/traits/v1/test_home.py

@@ -748,3 +748,142 @@ async def test_refresh_map_info_prefers_map_info_names_and_adds_missing_rooms(
     assert sorted_rooms[4].segment_id == 20
     assert sorted_rooms[4].name == "Office from rooms_trait"
     assert sorted_rooms[4].iot_id == "9999001"
+
+
+async def test_home_trait_listener_notifications(
+    home_trait: HomeTrait,
+    mock_rpc_channel: AsyncMock,
+    mock_mqtt_rpc_channel: AsyncMock,
+    mock_map_rpc_channel: AsyncMock,
+    device_cache: DeviceCache,
+    web_api_client: AsyncMock,
+) -> None:
+    """Test that listener callbacks are called when home discovery/updates occur."""
+    mock_callback = MagicMock()
+    unsub = home_trait.add_update_listener(mock_callback)
+    # Callback should NOT be executed immediately on subscription since home_map_info is None
+    assert mock_callback.call_count == 0
+
+    # 1. Test notification on empty cache discovery
+    mock_rpc_channel.send_command.side_effect = [
+        UPDATED_STATUS_MAP_123,
+        ROOM_MAPPING_DATA_MAP_123,
+        UPDATED_STATUS_MAP_0,
+        ROOM_MAPPING_DATA_MAP_0,
+    ]
+    mock_mqtt_rpc_channel.send_command.side_effect = [
+        MULTI_MAP_LIST_DATA,
+        {},
+        {},
+    ]
+    mock_map_rpc_channel.send_command.side_effect = [
+        MAP_BYTES_RESPONSE_2,
+        MAP_BYTES_RESPONSE_1,
+    ]
+    web_api_client.get_rooms.return_value = [
+        HomeDataRoom(id=2362048, name="Example room 1"),
+        HomeDataRoom(id=2362044, name="Example room 2"),
+    ]
+
+    await home_trait.discover_home()
+    # Mock callback should have been called during discovery
+    assert mock_callback.call_count > 0
+    mock_callback.reset_mock()
+
+    # 2. Test that registering a new listener now (when cache is populated) executes immediately
+    mock_callback_immediate = MagicMock()
+    unsub_immediate = home_trait.add_update_listener(mock_callback_immediate)
+    assert mock_callback_immediate.call_count == 1
+    unsub_immediate()
+
+    # 3. Test notification on cached discovery
+    # Re-run discover_home (which skips API calls and loads from cache)
+    await home_trait.discover_home()
+    assert mock_callback.call_count == 1
+    mock_callback.reset_mock()
+
+    # 4. Test notification on refresh
+    mock_rpc_channel.send_command.side_effect = [
+        ROOM_MAPPING_DATA_MAP_0,
+    ]
+    mock_mqtt_rpc_channel.send_command.side_effect = [
+        MULTI_MAP_LIST_DATA,
+    ]
+    mock_map_rpc_channel.send_command.side_effect = [
+        MAP_BYTES_RESPONSE_1,
+    ]
+    await home_trait.refresh()
+    assert mock_callback.call_count > 0
+
+    # Unsubscribe and verify it is no longer called
+    mock_callback.reset_mock()
+    unsub()
+    mock_rpc_channel.send_command.side_effect = [
+        ROOM_MAPPING_DATA_MAP_0,
+    ]
+    mock_mqtt_rpc_channel.send_command.side_effect = [
+        MULTI_MAP_LIST_DATA,
+    ]
+    mock_map_rpc_channel.send_command.side_effect = [
+        MAP_BYTES_RESPONSE_1,
+    ]
+    await home_trait.refresh()
+    assert mock_callback.call_count == 0
+
+
+async def test_home_trait_map_eviction(
+    home_trait: HomeTrait,
+    mock_rpc_channel: AsyncMock,
+    mock_mqtt_rpc_channel: AsyncMock,
+    mock_map_rpc_channel: AsyncMock,
+    device_cache: DeviceCache,
+    web_api_client: AsyncMock,
+) -> None:
+    """Test that maps deleted from the device are evicted from cache during refresh."""
+    # Pre-populate cache with maps 0 and 123
+    device_cache_data = DeviceCacheData(
+        home_map_info={
+            0: CombinedMapInfo(map_flag=0, name="Ground Floor", rooms=[]),
+            123: CombinedMapInfo(map_flag=123, name="Second Floor", rooms=[]),
+        },
+        home_map_content_base64={
+            0: base64.b64encode(MAP_BYTES_RESPONSE_1).decode("utf-8"),
+            123: base64.b64encode(MAP_BYTES_RESPONSE_2).decode("utf-8"),
+        },
+    )
+    await device_cache.set(device_cache_data)
+    await home_trait.discover_home()
+
+    assert home_trait.home_map_info is not None
+    assert home_trait.home_map_content is not None
+    assert len(home_trait.home_map_info) == 2
+    assert len(home_trait.home_map_content) == 2
+
+    # Set up listener callback
+    mock_callback = MagicMock()
+    home_trait.add_update_listener(mock_callback)
+    assert mock_callback.call_count == 1
+    mock_callback.reset_mock()
+
+    # Mock maps_trait.refresh so that only map 0 is returned (map 123 deleted)
+    mock_mqtt_rpc_channel.send_command.side_effect = [
+        MULTI_MAP_LIST_SINGLE_MAP_DATA,  # Only has map 0
+    ]
+    mock_rpc_channel.send_command.side_effect = [
+        ROOM_MAPPING_DATA_MAP_0,
+    ]
+    mock_map_rpc_channel.send_command.side_effect = [
+        MAP_BYTES_RESPONSE_1,
+    ]
+
+    await home_trait.refresh()
+
+    # Verify map 123 is excluded from memory cache
+    assert home_trait.home_map_info is not None
+    assert home_trait.home_map_content is not None
+    assert 123 not in home_trait.home_map_info
+    assert 123 not in home_trait.home_map_content
+    assert 0 in home_trait.home_map_info
+
+    # Verify listener notified
+    assert mock_callback.call_count > 0