fix: Improve Home trait discovery process.

This updates the Home trait to allow the current map to be refreshed even if initial discovery fails. The caller can invoke discovery at startup, but then after just invoke refresh and the right thing should happen.

We now track discovery explicitly and only persist the cache when we have completed discovery to avoid persisting an incomplete view of all available maps.

Author: allenporter

roborock/devices/traits/v1/home.py

@@ -7,6 +7,12 @@
 This trait depends on the MapsTrait and RoomsTrait to gather the necessary
 information. It provides properties to access the current map, the list of
 rooms with names, and the map image and data.
+
+Callers may first call `discover_home()` to populate the home layout cache by
+iterating through all available maps on the device. This will cache the map
+information and room names for all maps to minimize map switching and improve
+performance. After the initial discovery, callers can call `refresh()` to update
+the current map's information and room names as needed.
 """
 
 import asyncio
@@ -64,6 +70,7 @@ def __init__(
         self._map_content = map_content
         self._rooms_trait = rooms_trait
         self._cache = cache
+        self._discovery_completed = False
         self._home_map_info: dict[int, CombinedMapInfo] | None = None
         self._home_map_content: dict[int, MapContent] | None = None
 
@@ -82,6 +89,7 @@ async def discover_home(self) -> None:
         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
+            self._discovery_completed = True
             try:
                 self._home_map_content = {
                     k: self._map_content.parse_map_content(v) for k, v in cache_data.home_map_content.items()
@@ -101,6 +109,7 @@ async def discover_home(self) -> None:
 
         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))
+        self._discovery_completed = True
         await self._update_home_cache(home_map_info, home_map_content)
 
     async def _refresh_map_info(self, map_info) -> CombinedMapInfo:
@@ -156,8 +165,15 @@ async def refresh(self) -> Self:
         active maps or re-discover the home. It is expected that this will keep
         information up to date for the current map as users switch to that map.
         """
-        if self._home_map_info is None:
-            raise RoborockException("Cannot refresh home data without home cache, did you call discover_home()?")
+        if not self._discovery_completed:
+            # Running initial discovery also populates all of the same information
+            # as below so we can just call that method. If the device is busy
+            # then we'll fall through below to refresh the current map only.
+            try:
+                await self.discover_home()
+                return self
+            except RoborockDeviceBusy:
+                _LOGGER.debug("Cannot refresh home data while device is busy cleaning")
 
         # Refresh the list of map names/info
         await self._maps_trait.refresh()
@@ -170,7 +186,9 @@ async def refresh(self) -> Self:
         new_map_content = await self._refresh_map_content()
         # Refresh the current map's room 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)
+        await self._update_current_map(
+            map_flag, combined_map_info, new_map_content, update_cache=self._discovery_completed
+        )
         return self
 
     @property
@@ -206,16 +224,28 @@ async def _update_home_cache(
         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
+    async def _update_current_map(
+        self,
+        map_flag: int,
+        map_info: CombinedMapInfo,
+        map_content: MapContent,
+        update_cache: bool,
     ) -> 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")
+        # Update the persistent cache if requested e.g. home discovery has
+        # completed and we want to keep it fresh. Otherwise just update the
+        # in memory map below.
+        if update_cache:
+            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:
+            self._home_map_info = {}
         self._home_map_info[map_flag] = map_info
+
+        if self._home_map_content is None:
+            self._home_map_content = {}
         self._home_map_content[map_flag] = map_content

tests/devices/traits/v1/test_home.py

@@ -227,6 +227,11 @@ async def test_discover_home_empty_cache(
     assert current_map_data.map_flag == 0
     assert current_map_data.name == "Ground Floor"
 
+    # Verify the persistent cache has been updated
+    cache_data = await home_trait._cache.get()
+    assert len(cache_data.home_map_info) == 2
+    assert len(cache_data.home_map_content) == 2
+
 
 async def test_discover_home_with_existing_cache(
     home_trait: HomeTrait,
@@ -372,17 +377,6 @@ async def test_refresh_updates_current_map_cache(
     assert home_trait.home_map_content[0].image_content == TEST_IMAGE_CONTENT_1
 
 
-async def test_refresh_no_cache_no_update(
-    home_trait: HomeTrait,
-    mock_rpc_channel: AsyncMock,
-    mock_mqtt_rpc_channel: AsyncMock,
-) -> None:
-    """Test that refresh doesn't update when no cache exists."""
-    # 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()
-
-
 async def test_current_map_data_property(
     home_trait: HomeTrait,
     mock_rpc_channel: AsyncMock,
@@ -425,8 +419,14 @@ async def test_discover_home_device_busy_cleaning(
     home_trait: HomeTrait,
     mock_rpc_channel: AsyncMock,
     mock_mqtt_rpc_channel: AsyncMock,
+    mock_map_rpc_channel: AsyncMock,
 ) -> None:
-    """Test that discovery raises RoborockDeviceBusy when device is cleaning."""
+    """Test that discovery raises RoborockDeviceBusy when device is cleaning.
+
+    This tests the initial failure scenario during discovery where the device
+    is busy, then a retry attepmt where the device is still busy, then finally
+    a successful attempt to run discovery when the device is idle.
+    """
     # Set the status trait state to cleaning
     status_trait.state = RoborockStateCode.cleaning
 
@@ -438,6 +438,80 @@ async def test_discover_home_device_busy_cleaning(
     assert mock_rpc_channel.send_command.call_count == 0
     assert mock_mqtt_rpc_channel.send_command.call_count == 0
 
+    # Setup mocks for refresh
+    mock_rpc_channel.send_command.side_effect = [
+        ROOM_MAPPING_DATA_MAP_0,  # Room mapping refresh
+    ]
+    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
+    ]
+
+    # Now attempt to refresh the device while cleaning. This should still fail
+    # home discovery but allow refresh to update the current map.
+    await home_trait.refresh()
+
+    # Verify the home information is now populated
+    current_data = home_trait.current_map_data
+    assert current_data is not None
+    assert current_data.map_flag == 0
+    assert current_data.name == "Ground Floor"
+    assert home_trait.home_map_info is not None
+    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}
+    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
+
+    # Verify the persistent cache has not been updated since discovery
+    # has not fully completed.
+    cache_data = await home_trait._cache.get()
+    assert not cache_data.home_map_info
+    assert not cache_data.home_map_content
+
+    # Set the status trait state to idle which will mean we can attempt discovery
+    # on the next refresh. This should have the result of updating the
+    # persistent cache which is verified below.
+    status_trait.state = RoborockStateCode.idle
+
+    # Setup mocks for the discovery process
+    mock_rpc_channel.send_command.side_effect = [
+        UPDATED_STATUS_MAP_123,  # Status after switching to map 123
+        ROOM_MAPPING_DATA_MAP_123,  # Rooms for map 123
+        UPDATED_STATUS_MAP_0,  # Status after switching back to map 0
+        ROOM_MAPPING_DATA_MAP_0,  # Rooms for map 0
+    ]
+    mock_mqtt_rpc_channel.send_command.side_effect = [
+        MULTI_MAP_LIST_DATA,  # Multi maps list
+        {},  # 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
+    ]
+
+    # Refreshing should now perform discovery successfully
+    await home_trait.refresh()
+
+    # Verify we now have all of the information populated from discovery
+    assert home_trait.home_map_info is not None
+    assert len(home_trait.home_map_info) == 2
+
+    assert home_trait.home_map_info is not None
+    assert home_trait.home_map_info.keys() == {0, 123}
+    assert home_trait.home_map_content is not None
+    assert home_trait.home_map_content.keys() == {0, 123}
+
+    # Verify the persistent cache has been updated
+    cache_data = await home_trait._cache.get()
+    assert len(cache_data.home_map_info) == 2
+    assert len(cache_data.home_map_content) == 2
+
 
 async def test_single_map_no_switching(
     home_trait: HomeTrait,