Handle battery going down mid clean

Author: Lash-L

src/roborock_local_server/bundled_backend/shared/routine_runner.py

@@ -51,6 +51,7 @@ def _ensure_local_python_roborock_on_path() -> None:
 _STEP_START_POLL_INTERVAL_SECONDS = 0.5
 _STATUS_POLL_INTERVAL_SECONDS = 5.0
 _ROUTINE_READY_STATES = {3, 8, 100}
+_RESUME_BATTERY_THRESHOLD = 80
 _POST_STEP_SETTLE_SECONDS = 15.0
 _POST_STEP_SETTLE_TIMEOUT_SECONDS = 10 * 60
 from .inventory_io import WEB_API_INVENTORY_FILE
@@ -59,6 +60,11 @@ def _ensure_local_python_roborock_on_path() -> None:
     "do_scenes_segments",
     "do_scenes_zones",
 }
+_RESUME_COMMAND_BY_IN_CLEANING: dict[int, RoborockCommand] = {
+    RoborockInCleaning.global_clean_not_complete.value: RoborockCommand.APP_START,
+    RoborockInCleaning.zone_clean_not_complete.value: RoborockCommand.RESUME_ZONED_CLEAN,
+    RoborockInCleaning.segment_clean_not_complete.value: RoborockCommand.RESUME_SEGMENT_CLEAN,
+}
 
 
 class RoutineExecutionError(RuntimeError):
@@ -553,6 +559,7 @@ async def wait_for_step_complete(self) -> None:
         start_deadline = loop.time() + _STEP_START_TIMEOUT_SECONDS
         saw_activity = False
         saw_cleaning = False
+        sent_resume = False
 
         try:
             while True:
@@ -563,6 +570,7 @@ async def wait_for_step_complete(self) -> None:
                 status = await asyncio.wait_for(self.get_status(), timeout=remaining)
                 state = _enum_or_int_value(status.state)
                 in_cleaning = _enum_or_int_value(status.in_cleaning)
+                battery = int(status.battery) if status.battery is not None else -1
                 observed = (state, in_cleaning)
                 if observed != last_observed:
                     self._logger.info(
@@ -579,8 +587,35 @@ async def wait_for_step_complete(self) -> None:
                     in_cleaning == RoborockInCleaning.complete.value
                     and state in _ROUTINE_READY_STATES
                 )
+
+                charging_mid_clean = (
+                    state == 8
+                    and in_cleaning != RoborockInCleaning.complete.value
+                    and saw_cleaning
+                    and not sent_resume
+                    and battery >= _RESUME_BATTERY_THRESHOLD
+                )
+                if charging_mid_clean:
+                    resume_cmd = _RESUME_COMMAND_BY_IN_CLEANING.get(in_cleaning)
+                    if resume_cmd is not None:
+                        self._logger.info(
+                            "Routine wait: robot charging mid-clean (battery=%s%%), "
+                            "sending %s",
+                            battery,
+                            resume_cmd.value,
+                        )
+                        try:
+                            await self.send_command(resume_cmd, [])
+                            sent_resume = True
+                        except Exception as exc:  # noqa: BLE001
+                            self._logger.warning(
+                                "Routine wait: resume command failed: %s", exc
+                            )
+
                 if not is_ready:
                     saw_activity = True
+                    if sent_resume and state != 8:
+                        sent_resume = False
                 elif saw_cleaning:
                     return
                 elif saw_activity:

tests/test_routine_runner.py

@@ -276,6 +276,7 @@ def __init__(self, status_sequence: list[dict]) -> None:
         self._statuses = [StatusV2.from_dict(s) for s in status_sequence]
         self._index = 0
         self._logger = logging.getLogger("test-wait")
+        self.sent_commands: list[tuple[RoborockCommand, list | dict | None]] = []
 
     async def get_status(self) -> StatusV2:
         if self._index < len(self._statuses):
@@ -284,6 +285,9 @@ async def get_status(self) -> StatusV2:
             return status
         return self._statuses[-1]
 
+    async def send_command(self, command: RoborockCommand, params=None) -> None:
+        self.sent_commands.append((command, params))
+
 
 _ScriptedStatusClient.wait_for_step_complete = (
     routine_runner_module._RoutineMqttClient.wait_for_step_complete
@@ -368,3 +372,91 @@ async def exercise() -> None:
             await client.wait_for_step_complete()
 
     asyncio.run(exercise())
+
+
+def test_wait_for_step_complete_resume_after_mid_clean_charge(monkeypatch) -> None:
+    """Robot returns to dock mid-clean with low battery, charges, gets resumed, completes."""
+    monkeypatch.setattr(routine_runner_module, "_STEP_START_POLL_INTERVAL_SECONDS", 0.0)
+    monkeypatch.setattr(routine_runner_module, "_STATUS_POLL_INTERVAL_SECONDS", 0.0)
+    monkeypatch.setattr(routine_runner_module, "_RESUME_BATTERY_THRESHOLD", 80)
+
+    async def exercise() -> None:
+        client = _ScriptedStatusClient([
+            {"state": 18, "in_cleaning": 3, "battery": 55},  # segment cleaning
+            {"state": 26, "in_cleaning": 3, "battery": 30},  # going to wash mop
+            {"state": 6, "in_cleaning": 3, "battery": 14},   # returning home (low battery)
+            {"state": 8, "in_cleaning": 3, "battery": 14},   # charging mid-clean, too low
+            {"state": 8, "in_cleaning": 3, "battery": 50},   # still too low
+            {"state": 8, "in_cleaning": 3, "battery": 80},   # threshold reached → resume sent
+            {"state": 18, "in_cleaning": 3, "battery": 80},  # cleaning resumes
+            {"state": 18, "in_cleaning": 3, "battery": 60},
+            {"state": 6, "in_cleaning": 3, "battery": 40},   # returning home
+            {"state": 8, "in_cleaning": 0, "battery": 40},   # step complete
+        ])
+        await client.wait_for_step_complete()
+        assert len(client.sent_commands) == 1
+        assert client.sent_commands[0] == (RoborockCommand.RESUME_SEGMENT_CLEAN, [])
+
+    asyncio.run(exercise())
+
+
+def test_wait_for_step_complete_resume_zoned_clean(monkeypatch) -> None:
+    """Resume uses correct command for zone cleaning."""
+    monkeypatch.setattr(routine_runner_module, "_STEP_START_POLL_INTERVAL_SECONDS", 0.0)
+    monkeypatch.setattr(routine_runner_module, "_STATUS_POLL_INTERVAL_SECONDS", 0.0)
+    monkeypatch.setattr(routine_runner_module, "_RESUME_BATTERY_THRESHOLD", 80)
+
+    async def exercise() -> None:
+        client = _ScriptedStatusClient([
+            {"state": 18, "in_cleaning": 2, "battery": 50},  # zone cleaning
+            {"state": 8, "in_cleaning": 2, "battery": 80},   # charging mid-clean → resume
+            {"state": 18, "in_cleaning": 2, "battery": 80},  # resumes
+            {"state": 8, "in_cleaning": 0, "battery": 60},   # step complete
+        ])
+        await client.wait_for_step_complete()
+        assert len(client.sent_commands) == 1
+        assert client.sent_commands[0] == (RoborockCommand.RESUME_ZONED_CLEAN, [])
+
+    asyncio.run(exercise())
+
+
+def test_wait_for_step_complete_no_resume_when_battery_low(monkeypatch) -> None:
+    """No resume sent while battery is below threshold."""
+    monkeypatch.setattr(routine_runner_module, "_STEP_START_POLL_INTERVAL_SECONDS", 0.0)
+    monkeypatch.setattr(routine_runner_module, "_STATUS_POLL_INTERVAL_SECONDS", 0.0)
+    monkeypatch.setattr(routine_runner_module, "_STEP_COMPLETE_TIMEOUT_SECONDS", 0.1)
+    monkeypatch.setattr(routine_runner_module, "_RESUME_BATTERY_THRESHOLD", 80)
+
+    async def exercise() -> None:
+        client = _ScriptedStatusClient([
+            {"state": 18, "in_cleaning": 3, "battery": 50},  # cleaning
+            {"state": 8, "in_cleaning": 3, "battery": 14},   # charging, below threshold
+            {"state": 8, "in_cleaning": 3, "battery": 50},   # still below
+            {"state": 8, "in_cleaning": 3, "battery": 79},   # still below
+        ])
+        with pytest.raises(routine_runner_module.RoutineExecutionError, match="Timed out"):
+            await client.wait_for_step_complete()
+        assert len(client.sent_commands) == 0
+
+    asyncio.run(exercise())
+
+
+def test_wait_for_step_complete_resume_only_sent_once(monkeypatch) -> None:
+    """Resume command is only sent once even if robot returns to dock again."""
+    monkeypatch.setattr(routine_runner_module, "_STEP_START_POLL_INTERVAL_SECONDS", 0.0)
+    monkeypatch.setattr(routine_runner_module, "_STATUS_POLL_INTERVAL_SECONDS", 0.0)
+    monkeypatch.setattr(routine_runner_module, "_RESUME_BATTERY_THRESHOLD", 80)
+
+    async def exercise() -> None:
+        client = _ScriptedStatusClient([
+            {"state": 18, "in_cleaning": 3, "battery": 95},  # cleaning
+            {"state": 8, "in_cleaning": 3, "battery": 80},   # charging → resume sent
+            {"state": 18, "in_cleaning": 3, "battery": 80},  # cleaning resumes
+            {"state": 8, "in_cleaning": 3, "battery": 80},   # back to charging again → second resume sent
+            {"state": 18, "in_cleaning": 3, "battery": 80},  # cleaning resumes again
+            {"state": 8, "in_cleaning": 0, "battery": 60},   # step complete
+        ])
+        await client.wait_for_step_complete()
+        assert len(client.sent_commands) == 2
+
+    asyncio.run(exercise())