diff --git a/roborock/devices/cache.py b/roborock/devices/cache.py index 394db9e4..7f82abaa 100644 --- a/roborock/devices/cache.py +++ b/roborock/devices/cache.py @@ -25,6 +25,9 @@ class CacheData: home_map_info: dict[int, CombinedMapInfo] = field(default_factory=dict) """Home map information indexed by map_flag.""" + home_map_content: dict[int, bytes] = field(default_factory=dict) + """Home cache content for each map data indexed by map_flag.""" + device_features: DeviceFeatures | None = None """Device features information.""" diff --git a/roborock/devices/traits/v1/__init__.py b/roborock/devices/traits/v1/__init__.py index a5735511..0fc35bbf 100644 --- a/roborock/devices/traits/v1/__init__.py +++ b/roborock/devices/traits/v1/__init__.py @@ -141,7 +141,7 @@ def __init__( self.rooms = RoomsTrait(home_data) self.maps = MapsTrait(self.status) self.map_content = MapContentTrait(map_parser_config) - self.home = HomeTrait(self.status, self.maps, self.rooms, cache) + self.home = HomeTrait(self.status, self.maps, self.map_content, self.rooms, cache) self.device_features = DeviceFeaturesTrait(product.product_nickname, cache) self.network_info = NetworkInfoTrait(device_uid, cache) diff --git a/roborock/devices/traits/v1/home.py b/roborock/devices/traits/v1/home.py index 1cbfda01..b33dd573 100644 --- a/roborock/devices/traits/v1/home.py +++ b/roborock/devices/traits/v1/home.py @@ -20,6 +20,7 @@ from roborock.exceptions import RoborockDeviceBusy, RoborockException from roborock.roborock_typing import RoborockCommand +from .map_content import MapContent, MapContentTrait from .maps import MapsTrait from .rooms import RoomsTrait from .status import StatusTrait @@ -38,6 +39,7 @@ def __init__( self, status_trait: StatusTrait, maps_trait: MapsTrait, + map_content: MapContentTrait, rooms_trait: RoomsTrait, cache: Cache, ) -> None: @@ -59,9 +61,11 @@ def __init__( super().__init__() self._status_trait = status_trait self._maps_trait = maps_trait + self._map_content = map_content self._rooms_trait = rooms_trait self._cache = cache self._home_map_info: dict[int, CombinedMapInfo] | None = None + self._home_map_content: dict[int, MapContent] | None = None async def discover_home(self) -> None: """Iterate through all maps to discover rooms and cache them. @@ -75,10 +79,18 @@ async def discover_home(self) -> None: After discovery, the home cache will be populated and can be accessed via the `home_map_info` property. """ cache_data = await self._cache.get() - if cache_data.home_map_info: + if cache_data.home_map_info and cache_data.home_map_content: _LOGGER.debug("Home cache already populated, skipping discovery") self._home_map_info = cache_data.home_map_info - return + try: + self._home_map_content = { + k: self._map_content.parse_map_content(v) for k, v in cache_data.home_map_content.items() + } + except (ValueError, RoborockException) as ex: + _LOGGER.warning("Failed to parse cached home map content, will re-discover: %s", ex) + self._home_map_content = {} + else: + return if self._status_trait.state == RoborockStateCode.cleaning: raise RoborockDeviceBusy("Cannot perform home discovery while the device is cleaning") @@ -87,9 +99,9 @@ async def discover_home(self) -> None: if self._maps_trait.current_map_info is None: raise RoborockException("Cannot perform home discovery without current map info") - home_map_info = await self._build_home_map_info() + home_map_info, home_map_content = await self._build_home_map_info() _LOGGER.debug("Home discovery complete, caching data for %d maps", len(home_map_info)) - await self._update_home_map_info(home_map_info) + await self._update_home_cache(home_map_info, home_map_content) async def _refresh_map_info(self, map_info) -> CombinedMapInfo: """Collect room data for a specific map and return CombinedMapInfo.""" @@ -100,9 +112,19 @@ async def _refresh_map_info(self, map_info) -> CombinedMapInfo: rooms=self._rooms_trait.rooms or [], ) - async def _build_home_map_info(self) -> dict[int, CombinedMapInfo]: - """Perform the actual discovery and caching of home data.""" + async def _refresh_map_content(self) -> MapContent: + """Refresh the map content trait to get the latest map data.""" + await self._map_content.refresh() + return MapContent( + image_content=self._map_content.image_content, + map_data=self._map_content.map_data, + raw_api_response=self._map_content.raw_api_response, + ) + + async def _build_home_map_info(self) -> tuple[dict[int, CombinedMapInfo], dict[int, MapContent]]: + """Perform the actual discovery and caching of home map info and content.""" home_map_info: dict[int, CombinedMapInfo] = {} + home_map_content: dict[int, MapContent] = {} # Sort map_info to process the current map last, reducing map switching. # False (non-original maps) sorts before True (original map). We ensure @@ -120,9 +142,12 @@ async def _build_home_map_info(self) -> dict[int, CombinedMapInfo]: await self._maps_trait.set_current_map(map_info.map_flag) await asyncio.sleep(MAP_SLEEP) - map_data = await self._refresh_map_info(map_info) - home_map_info[map_info.map_flag] = map_data - return home_map_info + map_content = await self._refresh_map_content() + home_map_content[map_info.map_flag] = map_content + + combined_map_info = await self._refresh_map_info(map_info) + home_map_info[map_info.map_flag] = combined_map_info + return home_map_info, home_map_content async def refresh(self) -> Self: """Refresh current map's underlying map and room data, updating cache as needed. @@ -141,13 +166,11 @@ async def refresh(self) -> Self: ) is None: raise RoborockException("Cannot refresh home data without current map info") + # Refresh the map content to ensure we have the latest image and object positions + new_map_content = await self._refresh_map_content() # Refresh the current map's room data - current_map_data = self._home_map_info.get(map_flag) - if current_map_data: - map_data = await self._refresh_map_info(current_map_info) - if map_data != current_map_data: - await self._update_home_map_info({**self._home_map_info, map_flag: map_data}) - + combined_map_info = await self._refresh_map_info(current_map_info) + await self._update_current_map_cache(map_flag, combined_map_info, new_map_content) return self @property @@ -163,12 +186,36 @@ def current_map_data(self) -> CombinedMapInfo | None: return None return self._home_map_info.get(current_map_flag) + @property + def home_map_content(self) -> dict[int, MapContent] | None: + """Returns the map content for all cached maps.""" + return self._home_map_content + def _parse_response(self, response: common.V1ResponseData) -> Self: """This trait does not parse responses directly.""" raise NotImplementedError("HomeTrait does not support direct command responses") - async def _update_home_map_info(self, home_map_info: dict[int, CombinedMapInfo]) -> None: + async def _update_home_cache( + self, home_map_info: dict[int, CombinedMapInfo], home_map_content: dict[int, MapContent] + ) -> None: + """Update the entire home cache with new map info and content.""" cache_data = await self._cache.get() cache_data.home_map_info = home_map_info + cache_data.home_map_content = {k: v.raw_api_response for k, v in home_map_content.items() if v.raw_api_response} await self._cache.set(cache_data) self._home_map_info = home_map_info + self._home_map_content = home_map_content + + async def _update_current_map_cache( + self, map_flag: int, map_info: CombinedMapInfo, map_content: MapContent + ) -> None: + """Update the cache for the current map only.""" + cache_data = await self._cache.get() + cache_data.home_map_info[map_flag] = map_info + if map_content.raw_api_response: + cache_data.home_map_content[map_flag] = map_content.raw_api_response + await self._cache.set(cache_data) + if self._home_map_info is None or self._home_map_content is None: + raise RoborockException("Home cache is not initialized, cannot update current map cache") + self._home_map_info[map_flag] = map_info + self._home_map_content[map_flag] = map_content diff --git a/roborock/devices/traits/v1/map_content.py b/roborock/devices/traits/v1/map_content.py index 53e1aa3a..e0cc3e47 100644 --- a/roborock/devices/traits/v1/map_content.py +++ b/roborock/devices/traits/v1/map_content.py @@ -25,6 +25,13 @@ class MapContent(RoborockBase): map_data: MapData | None = None """The parsed map data which contains metadata for points on the map.""" + raw_api_response: bytes | None = None + """The raw bytes of the map data from the API for caching for future use. + + This should be treated as an opaque blob used only internally by this library + to re-parse the map data when needed. + """ + def __repr__(self) -> str: """Return a string representation of the MapContent.""" img = self.image_content @@ -48,7 +55,23 @@ def _parse_response(self, response: common.V1ResponseData) -> MapContent: """Parse the response from the device into a MapContentTrait instance.""" if not isinstance(response, bytes): raise ValueError(f"Unexpected MapContentTrait response format: {type(response)}") + return self.parse_map_content(response) + + def parse_map_content(self, response: bytes) -> MapContent: + """Parse the map content from raw bytes. + + This method is exposed so that cached map data can be parsed without + needing to go through the RPC channel. + + Args: + response: The raw bytes of the map data from the API. + + Returns: + MapContent: The parsed map content. + Raises: + RoborockException: If the map data cannot be parsed. + """ parsed_data = self._map_parser.parse(response) if parsed_data is None: raise ValueError("Failed to parse map data") @@ -56,4 +79,5 @@ def _parse_response(self, response: common.V1ResponseData) -> MapContent: return MapContent( image_content=parsed_data.image_content, map_data=parsed_data.map_data, + raw_api_response=response, ) diff --git a/tests/devices/traits/v1/test_home.py b/tests/devices/traits/v1/test_home.py index 9b96c285..27c6ce60 100644 --- a/tests/devices/traits/v1/test_home.py +++ b/tests/devices/traits/v1/test_home.py @@ -1,7 +1,7 @@ """Tests for the Home related functionality.""" from collections.abc import Iterator -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -10,10 +10,12 @@ from roborock.devices.cache import InMemoryCache from roborock.devices.device import RoborockDevice from roborock.devices.traits.v1.home import HomeTrait +from roborock.devices.traits.v1.map_content import MapContentTrait from roborock.devices.traits.v1.maps import MapsTrait from roborock.devices.traits.v1.rooms import RoomsTrait from roborock.devices.traits.v1.status import StatusTrait -from roborock.exceptions import RoborockDeviceBusy +from roborock.exceptions import RoborockDeviceBusy, RoborockException +from roborock.map.map_parser import ParsedMapData from roborock.roborock_typing import RoborockCommand from tests import mock_data @@ -40,6 +42,22 @@ ], } ] +MULTI_MAP_LIST_SINGLE_MAP_DATA = [ + { + "max_multi_map": 1, + "max_bak_map": 0, + "multi_map_count": 1, + "map_info": [ + { + "mapFlag": 0, + "add_time": 1747132930, + "length": 0, + "name": "Only Floor", + "bak_maps": [], + }, + ], + } +] ROOM_MAPPING_DATA_MAP_0 = [[16, "2362048"], [17, "2362044"]] ROOM_MAPPING_DATA_MAP_123 = [[18, "2362041"], [19, "2362042"]] @@ -53,6 +71,14 @@ **mock_data.STATUS, "map_status": 123 * 4 + 3, # Set current map to 123 } +MAP_BYTES_RESPONSE_1 = b"" +MAP_BYTES_RESPONSE_2 = b"" +TEST_IMAGE_CONTENT_1 = b"" +TEST_IMAGE_CONTENT_2 = b"" +TEST_PARSER_MAP = { + MAP_BYTES_RESPONSE_1: TEST_IMAGE_CONTENT_1, + MAP_BYTES_RESPONSE_2: TEST_IMAGE_CONTENT_2, +} @pytest.fixture(autouse=True) @@ -91,6 +117,13 @@ def maps_trait(device: RoborockDevice) -> MapsTrait: return device.v1_properties.maps +@pytest.fixture +def map_content_trait(device: RoborockDevice) -> MapContentTrait: + """Create a MapContentTrait instance with mocked dependencies.""" + assert device.v1_properties + return device.v1_properties.map_content + + @pytest.fixture def rooms_trait(device: RoborockDevice) -> RoomsTrait: """Create a RoomsTrait instance with mocked dependencies.""" @@ -99,9 +132,27 @@ def rooms_trait(device: RoborockDevice) -> RoomsTrait: @pytest.fixture -def home_trait(status_trait: StatusTrait, maps_trait: MapsTrait, rooms_trait: RoomsTrait, cache) -> HomeTrait: +def home_trait( + status_trait: StatusTrait, maps_trait: MapsTrait, map_content_trait: MapContentTrait, rooms_trait: RoomsTrait, cache +) -> HomeTrait: """Create a HomeTrait instance with mocked dependencies.""" - return HomeTrait(status_trait, maps_trait, rooms_trait, cache) + return HomeTrait(status_trait, maps_trait, map_content_trait, rooms_trait, cache) + + +@pytest.fixture(autouse=True) +def map_parser_fixture() -> Iterator[None]: + """Mock MapParser.parse to return predefined test map data.""" + + def parse_data(response: bytes) -> ParsedMapData: + if image_content := TEST_PARSER_MAP.get(response): + return ParsedMapData( + image_content=image_content, + map_data=MagicMock(), + ) + raise ValueError(f"Unexpected map bytes {response!r}") + + with patch("roborock.devices.traits.v1.map_content.MapParser.parse", side_effect=parse_data): + yield async def test_discover_home_empty_cache( @@ -109,6 +160,7 @@ async def test_discover_home_empty_cache( home_trait: HomeTrait, mock_rpc_channel: AsyncMock, mock_mqtt_rpc_channel: AsyncMock, + mock_map_rpc_channel: AsyncMock, ) -> None: """Test discovering home when cache is empty.""" # Setup mocks for the discovery process @@ -123,6 +175,10 @@ async def test_discover_home_empty_cache( {}, # LOAD_MULTI_MAP response for map 123 {}, # LOAD_MULTI_MAP response back to map 0 ] + mock_map_rpc_channel.send_command.side_effect = [ + MAP_BYTES_RESPONSE_2, # Map bytes for 123 + MAP_BYTES_RESPONSE_1, # Map bytes for 0 + ] # Before discovery, no cache should exist assert home_trait.home_map_info is None @@ -145,6 +201,11 @@ async def test_discover_home_empty_cache( assert map_0_data.rooms[1].segment_id == 17 assert map_0_data.rooms[1].name == "Example room 2" + map_0_content = home_trait.home_map_content[0] + assert map_0_content is not None + assert map_0_content.image_content == TEST_IMAGE_CONTENT_1 + assert map_0_content.map_data is not None + # Check map 123 data map_123_data = home_trait.home_map_info[123] assert map_123_data.map_flag == 123 @@ -155,6 +216,11 @@ async def test_discover_home_empty_cache( assert map_123_data.rooms[1].segment_id == 19 assert map_123_data.rooms[1].name == "Unknown" # Not in mock home data + map_123_content = home_trait.home_map_content[123] + assert map_123_content is not None + assert map_123_content.image_content == TEST_IMAGE_CONTENT_2 + assert map_123_content.map_data is not None + # Verify current map data is accessible current_map_data = home_trait.current_map_data assert current_map_data is not None @@ -171,6 +237,7 @@ async def test_discover_home_with_existing_cache( # Pre-populate the cache cache_data = await home_trait._cache.get() cache_data.home_map_info = {0: CombinedMapInfo(map_flag=0, name="Dummy", rooms=[])} + cache_data.home_map_content = {0: MAP_BYTES_RESPONSE_1} await home_trait._cache.set(cache_data) # Call discover_home @@ -182,6 +249,63 @@ async def test_discover_home_with_existing_cache( # Verify cache was loaded from storage assert home_trait.home_map_info == {0: CombinedMapInfo(map_flag=0, name="Dummy", rooms=[])} + assert home_trait.home_map_content + assert home_trait.home_map_content.keys() == {0} + map_0_content = home_trait.home_map_content[0] + assert map_0_content is not None + assert map_0_content.image_content == TEST_IMAGE_CONTENT_1 + assert map_0_content.map_data is not None + + +async def test_existing_home_cache_invalid_bytes( + home_trait: HomeTrait, + mock_rpc_channel: AsyncMock, + mock_mqtt_rpc_channel: AsyncMock, + mock_map_rpc_channel: AsyncMock, +) -> None: + """Test that discovery is skipped when cache already exists.""" + # Pre-populate the cache. + cache_data = await home_trait._cache.get() + cache_data.home_map_info = {0: CombinedMapInfo(map_flag=0, name="Dummy", rooms=[])} + # We override the map bytes parser to raise an exception above. + cache_data.home_map_content = {0: MAP_BYTES_RESPONSE_1} + await home_trait._cache.set(cache_data) + + # Setup mocks for the discovery process + mock_rpc_channel.send_command.side_effect = [ + ROOM_MAPPING_DATA_MAP_0, # Rooms for the single map + ] + mock_mqtt_rpc_channel.send_command.side_effect = [ + MULTI_MAP_LIST_SINGLE_MAP_DATA, # Single map list + ] + mock_map_rpc_channel.send_command.side_effect = [ + MAP_BYTES_RESPONSE_1, # Map bytes for the single map + ] + + # Call discover_home. First attempt raises an exception then loading from the server + # produes a valid result. + with patch( + "roborock.devices.traits.v1.map_content.MapParser.parse", + side_effect=[ + RoborockException("Invalid map bytes"), + ParsedMapData( + image_content=TEST_IMAGE_CONTENT_2, + map_data=MagicMock(), + ), + ], + ): + await home_trait.discover_home() + + # Verify cache was loaded from storage. The map was re-fetched from storage. + assert home_trait.home_map_info + assert home_trait.home_map_info.keys() == {0} + assert home_trait.home_map_info[0].name == "Only Floor" + assert home_trait.home_map_content + assert home_trait.home_map_content.keys() == {0} + map_0_content = home_trait.home_map_content[0] + assert map_0_content is not None + assert map_0_content.image_content == TEST_IMAGE_CONTENT_2 + assert map_0_content.map_data is not None async def test_discover_home_no_maps( @@ -200,18 +324,28 @@ async def test_discover_home_no_maps( async def test_refresh_updates_current_map_cache( + cache: InMemoryCache, status_trait: StatusTrait, home_trait: HomeTrait, mock_rpc_channel: AsyncMock, mock_mqtt_rpc_channel: AsyncMock, + mock_map_rpc_channel: AsyncMock, ) -> None: """Test that refresh updates the cache for the current map.""" # Pre-populate cache with some data - cache_data = await home_trait._cache.get() - + cache_data = await cache.get() cache_data.home_map_info = {0: CombinedMapInfo(map_flag=0, name="Old Ground Floor", rooms=[])} - await home_trait._cache.set(cache_data) - home_trait._home_map_info = cache_data.home_map_info + cache_data.home_map_content = {0: MAP_BYTES_RESPONSE_2} # Pre-existing different map bytes + await cache.set(cache_data) + await home_trait.discover_home() # Load cache into trait + # Verify initial cache state + assert home_trait.home_map_info + assert home_trait.home_map_info.keys() == {0} + assert home_trait.home_map_info[0].name == "Old Ground Floor" + assert len(home_trait.home_map_info[0].rooms) == 0 + assert home_trait.home_map_content + assert home_trait.home_map_content.keys() == {0} + assert home_trait.home_map_content[0].image_content == TEST_IMAGE_CONTENT_2 # Setup mocks for refresh mock_rpc_channel.send_command.side_effect = [ @@ -220,18 +354,22 @@ async def test_refresh_updates_current_map_cache( mock_mqtt_rpc_channel.send_command.side_effect = [ MULTI_MAP_LIST_DATA, # Maps refresh ] + mock_map_rpc_channel.send_command.side_effect = [ + MAP_BYTES_RESPONSE_1, # Map bytes refresh + ] # Perform refresh await home_trait.refresh() # Verify cache was updated for current map - updated_cache = home_trait.home_map_info - assert updated_cache is not None - assert 0 in updated_cache - - map_data = updated_cache[0] - assert map_data.name == "Ground Floor" # Updated from "Old Ground Floor" - assert len(map_data.rooms) == 2 + assert home_trait.home_map_info + assert home_trait.home_map_info.keys() == {0} + assert home_trait.home_map_info[0].name == "Ground Floor" + assert len(home_trait.home_map_info[0].rooms) == 2 + # Verify map content + assert home_trait.home_map_content + assert home_trait.home_map_content.keys() == {0} + assert home_trait.home_map_content[0].image_content == TEST_IMAGE_CONTENT_1 async def test_refresh_no_cache_no_update( @@ -240,10 +378,6 @@ async def test_refresh_no_cache_no_update( mock_mqtt_rpc_channel: AsyncMock, ) -> None: """Test that refresh doesn't update when no cache exists.""" - # Setup mocks for refresh - # mock_mqtt_rpc_channel.send_command.side_effect = [ - # MULTI_MAP_LIST_DATA, # Maps refresh - # ] # Perform refresh without existing cache with pytest.raises(Exception, match="Cannot refresh home data without home cache, did you call discover_home()?"): await home_trait.refresh() @@ -253,6 +387,7 @@ async def test_current_map_data_property( home_trait: HomeTrait, mock_rpc_channel: AsyncMock, mock_mqtt_rpc_channel: AsyncMock, + mock_map_rpc_channel: AsyncMock, ) -> None: """Test current_map_data property returns correct data.""" # Setup discovery @@ -267,6 +402,10 @@ async def test_current_map_data_property( {}, # LOAD_MULTI_MAP response for map 123 {}, # LOAD_MULTI_MAP response back to map 0 ] + mock_map_rpc_channel.send_command.side_effect = [ + MAP_BYTES_RESPONSE_2, # Map bytes for 123 + MAP_BYTES_RESPONSE_1, # Map bytes for 0 + ] await home_trait.discover_home() @@ -304,38 +443,26 @@ async def test_single_map_no_switching( home_trait: HomeTrait, mock_rpc_channel: AsyncMock, mock_mqtt_rpc_channel: AsyncMock, + mock_map_rpc_channel: AsyncMock, ) -> None: """Test that single map discovery doesn't trigger map switching.""" - single_map_data = [ - { - "max_multi_map": 1, - "max_bak_map": 0, - "multi_map_count": 1, - "map_info": [ - { - "mapFlag": 0, - "add_time": 1747132930, - "length": 0, - "name": "Only Floor", - "bak_maps": [], - }, - ], - } - ] - mock_rpc_channel.send_command.side_effect = [ ROOM_MAPPING_DATA_MAP_0, # Rooms for the single map ] mock_mqtt_rpc_channel.send_command.side_effect = [ - single_map_data, # Single map list + MULTI_MAP_LIST_SINGLE_MAP_DATA, # Single map list + ] + mock_map_rpc_channel.send_command.side_effect = [ + MAP_BYTES_RESPONSE_1, # Map bytes for the single map ] await home_trait.discover_home() # Verify cache is populated assert home_trait.home_map_info is not None - assert len(home_trait.home_map_info) == 1 - assert 0 in home_trait.home_map_info + assert home_trait.home_map_info.keys() == {0} + assert home_trait.home_map_content is not None + assert home_trait.home_map_content.keys() == {0} # Verify no LOAD_MULTI_MAP commands were sent (no map switching) load_map_calls = [