Fix bug where Roborock loading map in cleaning causes a crash (#153011)

homeassistant/components/roborock/coordinator.py

@@ -418,14 +418,36 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
             # If we don't have a cur map(shouldn't happen) just
             # return as we can't do anything.
             return
-        map_flags = sorted(self.maps, key=lambda data: data == cur_map, reverse=True)
+        if self.data.status.in_cleaning:
+            # If the vacuum is cleaning, we cannot change maps
+            # as it will interrupt the cleaning.
+            _LOGGER.info(
+                "Vacuum is cleaning, not switching to other maps to fetch rooms"
+            )
+            # Since this is hitting the cloud api, we want to be careful and will just
+            # stop here rather than retrying in the future.
+            map_flags = [cur_map]
+        else:
+            map_flags = sorted(
+                self.maps, key=lambda data: data == cur_map, reverse=True
+            )
         for map_flag in map_flags:
             if map_flag != cur_map:
                 # Only change the map and sleep if we have multiple maps.
-                await self.cloud_api.load_multi_map(map_flag)
-                self.current_map = map_flag
+                try:
+                    await self.cloud_api.load_multi_map(map_flag)
+                except RoborockException as ex:
+                    _LOGGER.debug(
+                        "Failed to change to map %s when refreshing maps: %s",
+                        map_flag,
+                        ex,
+                    )
+                    continue
+                else:
+                    self.current_map = map_flag
                 # We cannot get the map until the roborock servers fully process the
-                # map change.
+                # map change. If the above command fails, we should still sleep, just
+                # in case it executes delayed.
                 await asyncio.sleep(MAP_SLEEP)
             tasks = [self.set_current_map_rooms()]
             # The image is set within async_setup, so if it exists, we have it here.
@@ -436,11 +458,18 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
             # If either of these fail, we don't care, and we want to continue.
             await asyncio.gather(*tasks, return_exceptions=True)
 
-        if len(self.maps) > 1:
+        if len(self.maps) > 1 and not self.data.status.in_cleaning:
             # Set the map back to the map the user previously had selected so that it
             # does not change the end user's app.
             # Only needs to happen when we changed maps above.
-            await self.cloud_api.load_multi_map(cur_map)
+            try:
+                await self.cloud_api.load_multi_map(cur_map)
+            except RoborockException as ex:
+                _LOGGER.warning(
+                    "Failed to change back to map %s when refreshing maps: %s",
+                    cur_map,
+                    ex,
+                )
             self.current_map = cur_map
 
 

tests/components/roborock/test_coordinator.py

@@ -1,5 +1,6 @@
 """Test Roborock Coordinator specific logic."""
 
+import asyncio
 import copy
 from datetime import timedelta
 from unittest.mock import patch
@@ -11,13 +12,15 @@ from vacuum_map_parser_base.config.color import SupportedColor
 
 from homeassistant.components.roborock.const import (
     CONF_SHOW_BACKGROUND,
+    DOMAIN,
+    GET_MAPS_SERVICE_NAME,
     V1_CLOUD_IN_CLEANING_INTERVAL,
     V1_CLOUD_NOT_CLEANING_INTERVAL,
     V1_LOCAL_IN_CLEANING_INTERVAL,
     V1_LOCAL_NOT_CLEANING_INTERVAL,
 )
 from homeassistant.components.roborock.coordinator import RoborockDataUpdateCoordinator
-from homeassistant.const import Platform
+from homeassistant.const import ATTR_ENTITY_ID, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.util import dt as dt_util
 
@@ -29,7 +32,7 @@ from tests.common import MockConfigEntry, async_fire_time_changed
 @pytest.fixture
 def platforms() -> list[Platform]:
     """Fixture to set platforms used in the test."""
-    return [Platform.SENSOR]
+    return [Platform.SENSOR, Platform.VACUUM]
 
 
 @pytest.mark.parametrize(
@@ -166,3 +169,112 @@ async def test_no_maps(
     ):
         await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
     assert load_map.call_count == 0
+
+
+async def test_two_maps_in_cleaning(
+    hass: HomeAssistant,
+    mock_roborock_entry: MockConfigEntry,
+    bypass_api_fixture: None,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that we gracefully handle having two maps but we are in cleaning."""
+    prop = copy.deepcopy(PROP)
+    prop.status.in_cleaning = True
+    with (
+        patch(
+            "homeassistant.components.roborock.coordinator.RoborockLocalClientV1.get_prop",
+            return_value=prop,
+        ),
+        patch(
+            "homeassistant.components.roborock.RoborockMqttClientV1.load_multi_map"
+        ) as load_map,
+    ):
+        await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
+    # We should not try to load any maps as we should just get the information for our
+    # current map and move on.
+    assert load_map.call_count == 0
+    assert (
+        "Vacuum is cleaning, not switching to other maps to fetch rooms" in caplog.text
+    )
+
+
+async def test_failed_load_multi_map(
+    hass: HomeAssistant,
+    mock_roborock_entry: MockConfigEntry,
+    bypass_api_fixture: None,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that we gracefully handle one map failing to load."""
+    with (
+        patch(
+            "homeassistant.components.roborock.RoborockMqttClientV1.load_multi_map",
+            side_effect=[RoborockException(), None, None, None],
+        ) as load_map,
+    ):
+        await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
+    assert "Failed to change to map 1 when refreshing maps" in caplog.text
+    # We continue to try and load the next map so we we should have multiple load maps.
+    # 2 for both devices, even though one for one of the devices failed.
+    assert load_map.call_count == 4
+    # Just to be safe since we load the maps asynchronously, lets make sure that only
+    # one map out of the four didn't get called.
+    responses = await asyncio.gather(
+        *(
+            hass.services.async_call(
+                DOMAIN,
+                GET_MAPS_SERVICE_NAME,
+                {ATTR_ENTITY_ID: dev},
+                blocking=True,
+                return_response=True,
+            )
+            for dev in ("vacuum.roborock_s7_maxv", "vacuum.roborock_s7_2")
+        )
+    )
+    num_no_rooms = sum(
+        1
+        for res in responses
+        for data in res.values()
+        for m in data["maps"]
+        if not m["rooms"]
+    )
+    assert num_no_rooms == 1
+
+
+async def test_failed_reset_map(
+    hass: HomeAssistant,
+    mock_roborock_entry: MockConfigEntry,
+    bypass_api_fixture: None,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that we gracefully handle not being able to revert back to the original map."""
+    with (
+        patch(
+            "homeassistant.components.roborock.RoborockMqttClientV1.load_multi_map",
+            side_effect=[None, None, None, RoborockException()],
+        ) as load_map,
+    ):
+        await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
+    assert "Failed to change back to map 0 when refreshing maps" in caplog.text
+    # 2 for both devices, even though one for one of the devices failed.
+    assert load_map.call_count == 4
+    responses = await asyncio.gather(
+        *(
+            hass.services.async_call(
+                DOMAIN,
+                GET_MAPS_SERVICE_NAME,
+                {ATTR_ENTITY_ID: dev},
+                blocking=True,
+                return_response=True,
+            )
+            for dev in ("vacuum.roborock_s7_maxv", "vacuum.roborock_s7_2")
+        )
+    )
+    num_no_rooms = sum(
+        1
+        for res in responses
+        for data in res.values()
+        for m in data["maps"]
+        if not m["rooms"]
+    )
+    # No maps should be missing information, as we just couldn't go back to the original.
+    assert num_no_rooms == 0