diff --git a/roborock/devices/traits/v1/home.py b/roborock/devices/traits/v1/home.py index 2d9d7bb1..afc134ae 100644 --- a/roborock/devices/traits/v1/home.py +++ b/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() diff --git a/tests/devices/traits/v1/test_home.py b/tests/devices/traits/v1/test_home.py index 9c070e01..42660f6e 100644 --- a/tests/devices/traits/v1/test_home.py +++ b/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