fix: Don't perform discovery when the device is cleaning

Author: allenporter

roborock/devices/traits/v1/__init__.py

@@ -71,7 +71,7 @@ def __init__(
         self.rooms = RoomsTrait(home_data)
         self.maps = MapsTrait(self.status)
         self.map_content = MapContentTrait(map_parser_config)
-        self.home = HomeTrait(self.maps, self.rooms, cache)
+        self.home = HomeTrait(self.status, self.maps, self.rooms, cache)
         # This is a hack to allow setting the rpc_channel on all traits. This is
         # used so we can preserve the dataclass behavior when the values in the
         # traits are updated, but still want to allow them to have a reference

roborock/devices/traits/v1/home.py

@@ -13,14 +13,16 @@
 import logging
 from typing import Self
 
+from roborock.code_mappings import RoborockStateCode
 from roborock.containers import CombinedMapInfo, RoborockBase
 from roborock.devices.cache import Cache
 from roborock.devices.traits.v1 import common
-from roborock.exceptions import RoborockException
+from roborock.exceptions import RoborockDeviceBusy, RoborockException
 from roborock.roborock_typing import RoborockCommand
 
 from .maps import MapsTrait
 from .rooms import RoomsTrait
+from .status import StatusTrait
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -32,7 +34,13 @@ class HomeTrait(RoborockBase, common.V1TraitMixin):
 
     command = RoborockCommand.GET_MAP_V1  # This is not used
 
-    def __init__(self, maps_trait: MapsTrait, rooms_trait: RoomsTrait, cache: Cache) -> None:
+    def __init__(
+        self,
+        status_trait: StatusTrait,
+        maps_trait: MapsTrait,
+        rooms_trait: RoomsTrait,
+        cache: Cache,
+    ) -> None:
         """Initialize the HomeTrait.
 
         We keep track of the MapsTrait and RoomsTrait to provide a comprehensive
@@ -49,6 +57,7 @@ def __init__(self, maps_trait: MapsTrait, rooms_trait: RoomsTrait, cache: Cache)
         accuracy.
         """
         super().__init__()
+        self._status_trait = status_trait
         self._maps_trait = maps_trait
         self._rooms_trait = rooms_trait
         self._cache = cache
@@ -58,21 +67,29 @@ async def discover_home(self) -> None:
         """Iterate through all maps to discover rooms and cache them.
 
         This will be a no-op if the home cache is already populated.
+
+        This cannot be called while the device is cleaning, as that would interrupt the
+        cleaning process. This will raise `RoborockDeviceBusy` if the device is
+        currently cleaning.
+
+        After discovery, the home cache will be populated and can be accessed via the `home_cache` property.
         """
         cache_data = await self._cache.get()
         if cache_data.home_cache:
             _LOGGER.debug("Home cache already populated, skipping discovery")
             self._home_cache = cache_data.home_cache
             return
 
+        if self._status_trait.state == RoborockStateCode.cleaning:
+            raise RoborockDeviceBusy("Cannot perform home discovery while the device is cleaning")
+
         await self._maps_trait.refresh()
         if self._maps_trait.current_map_info is None:
             raise RoborockException("Cannot perform home discovery without current map info")
 
         home_cache = await self._build_home_cache()
         _LOGGER.debug("Home discovery complete, caching data for %d maps", len(home_cache))
         await self._update_home_cache(home_cache)
-        self._home_cache = home_cache
 
     async def _refresh_map_data(self, map_info) -> CombinedMapInfo:
         """Collect room data for a specific map and return CombinedMapInfo."""
@@ -129,8 +146,7 @@ async def refresh(self) -> Self:
         if current_map_data:
             map_data = await self._refresh_map_data(current_map_info)
             if map_data != current_map_data:
-                self._home_cache[map_flag] = map_data
-                await self._update_home_cache(self._home_cache)
+                await self._update_home_cache({**self._home_cache, map_flag: map_data})
 
         return self
 
@@ -155,3 +171,4 @@ async def _update_home_cache(self, home_cache: dict[int, CombinedMapInfo]) -> No
         cache_data = await self._cache.get()
         cache_data.home_cache = home_cache
         await self._cache.set(cache_data)
+        self._home_cache = home_cache

roborock/exceptions.py

@@ -81,3 +81,7 @@ class RoborockRateLimit(RoborockException):
 
 class RoborockNoResponseFromBaseURL(RoborockException):
     """We could not find an url that had a record of the given account."""
+
+
+class RoborockDeviceBusy(RoborockException):
+    """Class for Roborock device busy exceptions."""

tests/devices/traits/v1/test_home.py

@@ -5,13 +5,15 @@
 
 import pytest
 
+from roborock.code_mappings import RoborockStateCode
 from roborock.containers import CombinedMapInfo
 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.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.roborock_typing import RoborockCommand
 from tests import mock_data
 
@@ -97,9 +99,9 @@ def rooms_trait(device: RoborockDevice) -> RoomsTrait:
 
 
 @pytest.fixture
-def home_trait(maps_trait: MapsTrait, rooms_trait: RoomsTrait, cache) -> HomeTrait:
+def home_trait(status_trait: StatusTrait, maps_trait: MapsTrait, rooms_trait: RoomsTrait, cache) -> HomeTrait:
     """Create a HomeTrait instance with mocked dependencies."""
-    return HomeTrait(maps_trait, rooms_trait, cache)
+    return HomeTrait(status_trait, maps_trait, rooms_trait, cache)
 
 
 async def test_discover_home_empty_cache(
@@ -279,6 +281,25 @@ async def test_current_map_data_property(
     assert home_trait.current_map_data is None
 
 
+async def test_discover_home_device_busy_cleaning(
+    status_trait: StatusTrait,
+    home_trait: HomeTrait,
+    mock_rpc_channel: AsyncMock,
+    mock_mqtt_rpc_channel: AsyncMock,
+) -> None:
+    """Test that discovery raises RoborockDeviceBusy when device is cleaning."""
+    # Set the status trait state to cleaning
+    status_trait.state = RoborockStateCode.cleaning
+
+    # Attempt to discover home while cleaning
+    with pytest.raises(RoborockDeviceBusy, match="Cannot perform home discovery while the device is cleaning"):
+        await home_trait.discover_home()
+
+    # Verify no RPC calls were made (discovery was prevented)
+    assert mock_rpc_channel.send_command.call_count == 0
+    assert mock_mqtt_rpc_channel.send_command.call_count == 0
+
+
 async def test_single_map_no_switching(
     home_trait: HomeTrait,
     mock_rpc_channel: AsyncMock,