From 74f14204101995e5d9dd03bc2b566c78fdaefa7a Mon Sep 17 00:00:00 2001
From: Luke Lashley <conway220@gmail.com>
Date: Sun, 3 Mar 2024 12:41:31 -0500
Subject: [PATCH] Add the ability to resume cleaning on start button (#112122)

---
 homeassistant/components/roborock/vacuum.py |  7 +++-
 tests/components/roborock/test_vacuum.py    | 45 +++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/roborock/vacuum.py b/homeassistant/components/roborock/vacuum.py
index dafbb731bd2..3684b584047 100644
--- a/homeassistant/components/roborock/vacuum.py
+++ b/homeassistant/components/roborock/vacuum.py
@@ -121,7 +121,12 @@ class RoborockVacuum(RoborockCoordinatedEntity, StateVacuumEntity):
 
     async def async_start(self) -> None:
         """Start the vacuum."""
-        await self.send(RoborockCommand.APP_START)
+        if self._device_status.in_cleaning == 2:
+            await self.send(RoborockCommand.RESUME_ZONED_CLEAN)
+        elif self._device_status.in_cleaning == 3:
+            await self.send(RoborockCommand.RESUME_SEGMENT_CLEAN)
+        else:
+            await self.send(RoborockCommand.APP_START)
 
     async def async_pause(self) -> None:
         """Pause the vacuum."""
diff --git a/tests/components/roborock/test_vacuum.py b/tests/components/roborock/test_vacuum.py
index ecc501cc542..61a2ef5d8e2 100644
--- a/tests/components/roborock/test_vacuum.py
+++ b/tests/components/roborock/test_vacuum.py
@@ -1,6 +1,7 @@
 """Tests for Roborock vacuums."""
 
 
+import copy
 from typing import Any
 from unittest.mock import patch
 
@@ -8,6 +9,7 @@ import pytest
 from roborock import RoborockException
 from roborock.roborock_typing import RoborockCommand
 
+from homeassistant.components.roborock import DOMAIN
 from homeassistant.components.vacuum import (
     SERVICE_CLEAN_SPOT,
     SERVICE_LOCATE,
@@ -22,8 +24,10 @@ from homeassistant.const import ATTR_ENTITY_ID, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
+from homeassistant.setup import async_setup_component
 
 from tests.common import MockConfigEntry
+from tests.components.roborock.mock_data import PROP
 
 ENTITY_ID = "vacuum.roborock_s7_maxv"
 DEVICE_ID = "abc123"
@@ -90,6 +94,47 @@ async def test_commands(
         assert mock_send_command.call_args[0][1] == called_params
 
 
+@pytest.mark.parametrize(
+    ("in_cleaning_int", "expected_command"),
+    [
+        (0, RoborockCommand.APP_START),
+        (1, RoborockCommand.APP_START),
+        (2, RoborockCommand.RESUME_ZONED_CLEAN),
+        (3, RoborockCommand.RESUME_SEGMENT_CLEAN),
+    ],
+)
+async def test_resume_cleaning(
+    hass: HomeAssistant,
+    bypass_api_fixture,
+    mock_roborock_entry: MockConfigEntry,
+    in_cleaning_int: int,
+    expected_command: RoborockCommand,
+) -> None:
+    """Test resuming clean on start button when a clean is paused."""
+    prop = copy.deepcopy(PROP)
+    prop.status.in_cleaning = in_cleaning_int
+    with patch(
+        "homeassistant.components.roborock.coordinator.RoborockLocalClient.get_prop",
+        return_value=prop,
+    ):
+        await async_setup_component(hass, DOMAIN, {})
+    vacuum = hass.states.get(ENTITY_ID)
+    assert vacuum
+
+    data = {ATTR_ENTITY_ID: ENTITY_ID}
+    with patch(
+        "homeassistant.components.roborock.coordinator.RoborockLocalClient.send_command"
+    ) as mock_send_command:
+        await hass.services.async_call(
+            Platform.VACUUM,
+            SERVICE_START,
+            data,
+            blocking=True,
+        )
+        assert mock_send_command.call_count == 1
+        assert mock_send_command.call_args[0][0] == expected_command
+
+
 async def test_failed_user_command(
     hass: HomeAssistant,
     bypass_api_fixture,