From 6754bf2466e27840c3f7862ef14723bc1d57efe7 Mon Sep 17 00:00:00 2001
From: Luke Lashley <conway220@gmail.com>
Date: Sat, 8 Mar 2025 13:04:40 -0500
Subject: [PATCH] Send Roborock commands via cloud api when needed (#138496)

* Send via cloud api when needed

* Extract logic to helper function

* change to class method
---
 homeassistant/components/roborock/entity.py | 29 ++++++++++++++++-----
 tests/components/roborock/test_vacuum.py    | 24 +++++++++++++++++
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/roborock/entity.py b/homeassistant/components/roborock/entity.py
index 4a16ada5967..d417ac17159 100644
--- a/homeassistant/components/roborock/entity.py
+++ b/homeassistant/components/roborock/entity.py
@@ -8,7 +8,11 @@ from roborock.containers import Consumable, Status
 from roborock.exceptions import RoborockException
 from roborock.roborock_message import RoborockDataProtocol
 from roborock.roborock_typing import RoborockCommand
-from roborock.version_1_apis.roborock_client_v1 import AttributeCache, RoborockClientV1
+from roborock.version_1_apis.roborock_client_v1 import (
+    CLOUD_REQUIRED,
+    AttributeCache,
+    RoborockClientV1,
+)
 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
 from roborock.version_a01_apis import RoborockClientA01
 
@@ -53,14 +57,16 @@ class RoborockEntityV1(RoborockEntity):
         """Get an item from the api cache."""
         return self._api.cache[attribute]
 
-    async def send(
-        self,
+    @classmethod
+    async def _send_command(
+        cls,
         command: RoborockCommand | str,
+        api: RoborockClientV1,
         params: dict[str, Any] | list[Any] | int | None = None,
     ) -> dict:
-        """Send a command to a vacuum cleaner."""
+        """Send a Roborock command with params to a given api."""
         try:
-            response: dict = await self._api.send_command(command, params)
+            response: dict = await api.send_command(command, params)
         except RoborockException as err:
             if isinstance(command, RoborockCommand):
                 command_name = command.name
@@ -75,6 +81,14 @@ class RoborockEntityV1(RoborockEntity):
             ) from err
         return response
 
+    async def send(
+        self,
+        command: RoborockCommand | str,
+        params: dict[str, Any] | list[Any] | int | None = None,
+    ) -> dict:
+        """Send a command to a vacuum cleaner."""
+        return await self._send_command(command, self._api, params)
+
     @property
     def api(self) -> RoborockClientV1:
         """Returns the api."""
@@ -152,7 +166,10 @@ class RoborockCoordinatedEntityV1(
         params: dict[str, Any] | list[Any] | int | None = None,
     ) -> dict:
         """Overloads normal send command but refreshes coordinator."""
-        res = await super().send(command, params)
+        if command in CLOUD_REQUIRED:
+            res = await self._send_command(command, self.coordinator.cloud_api, params)
+        else:
+            res = await self._send_command(command, self._api, params)
         await self.coordinator.async_refresh()
         return res
 
diff --git a/tests/components/roborock/test_vacuum.py b/tests/components/roborock/test_vacuum.py
index d9d4340ec83..15fdeb4767c 100644
--- a/tests/components/roborock/test_vacuum.py
+++ b/tests/components/roborock/test_vacuum.py
@@ -117,6 +117,30 @@ async def test_commands(
         assert mock_send_command.call_args[0][1] == called_params
 
 
+async def test_cloud_command(
+    hass: HomeAssistant,
+    bypass_api_fixture,
+    setup_entry: MockConfigEntry,
+) -> None:
+    """Test sending commands to the vacuum."""
+
+    vacuum = hass.states.get(ENTITY_ID)
+    assert vacuum
+
+    data = {ATTR_ENTITY_ID: ENTITY_ID, "command": "get_map_v1"}
+    with patch(
+        "homeassistant.components.roborock.coordinator.RoborockMqttClientV1.send_command"
+    ) as mock_send_command:
+        await hass.services.async_call(
+            Platform.VACUUM,
+            SERVICE_SEND_COMMAND,
+            data,
+            blocking=True,
+        )
+        assert mock_send_command.call_count == 1
+        assert mock_send_command.call_args[0][0] == RoborockCommand.GET_MAP_V1
+
+
 @pytest.mark.parametrize(
     ("in_cleaning_int", "expected_command"),
     [