Add scene support to roborock (#137203)

* feature: add scene buttons to roborock * feature: upgrade python-roborock * feature: upgrade python-roborock * feature: upgrade python-roborock * feature: upgrade python-roborock * feature: upgrade python-roborock * feature: upgrade python-roborock * feature: upgrade python-roborock * feature: upgrade python-roborock * feature: upgrade python-roborock
2025-02-09 23:39:38 +02:00 · 2025-02-09 23:39:38 +02:00 · 379bf10675
parent e8e4d2a83c
commit 379bf10675
7 changed files with 278 additions and 12 deletions
--- a/homeassistant/components/roborock/init.py
+++ b/homeassistant/components/roborock/init.py
@ -82,7 +82,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) ->
    # Get a Coordinator if the device is available or if we have connected to the device before
    coordinators = await asyncio.gather(
        *build_setup_functions(
-            hass, entry, device_map, user_data, product_info, home_data.rooms
+            hass,
+            entry,
+            device_map,
+            user_data,
+            product_info,
+            home_data.rooms,
+            api_client,
        ),
        return_exceptions=True,
    )
@ -134,6 +140,7 @@ def build_setup_functions(
    user_data: UserData,
    product_info: dict[str, HomeDataProduct],
    home_data_rooms: list[HomeDataRoom],
+    api_client: RoborockApiClient,
 ) -> list[
    Coroutine[
        Any,
@ -150,6 +157,7 @@ def build_setup_functions(
            device,
            product_info[device.product_id],
            home_data_rooms,
+            api_client,
        )
        for device in device_map.values()
    ]
@ -162,11 +170,12 @@ async def setup_device(
    device: HomeDataDevice,
    product_info: HomeDataProduct,
    home_data_rooms: list[HomeDataRoom],
+    api_client: RoborockApiClient,
 ) -> RoborockDataUpdateCoordinator | RoborockDataUpdateCoordinatorA01 | None:
    """Set up a coordinator for a given device."""
    if device.pv == "1.0":
        return await setup_device_v1(
-            hass, entry, user_data, device, product_info, home_data_rooms
+            hass, entry, user_data, device, product_info, home_data_rooms, api_client
        )
    if device.pv == "A01":
        return await setup_device_a01(hass, entry, user_data, device, product_info)
@ -186,6 +195,7 @@ async def setup_device_v1(
    device: HomeDataDevice,
    product_info: HomeDataProduct,
    home_data_rooms: list[HomeDataRoom],
+    api_client: RoborockApiClient,
 ) -> RoborockDataUpdateCoordinator | None:
    """Set up a device Coordinator."""
    mqtt_client = await hass.async_add_executor_job(
@ -207,7 +217,15 @@ async def setup_device_v1(
        await mqtt_client.async_release()
        raise
    coordinator = RoborockDataUpdateCoordinator(
-        hass, entry, device, networking, product_info, mqtt_client, home_data_rooms
+        hass,
+        entry,
+        device,
+        networking,
+        product_info,
+        mqtt_client,
+        home_data_rooms,
+        api_client,
+        user_data,
    )
    try:
        await coordinator.async_config_entry_first_refresh()
--- a/homeassistant/components/roborock/const.py
+++ b/homeassistant/components/roborock/const.py
@ -36,6 +36,7 @@ PLATFORMS = [
    Platform.BUTTON,
    Platform.IMAGE,
    Platform.NUMBER,
+    Platform.SCENE,
    Platform.SELECT,
    Platform.SENSOR,
    Platform.SWITCH,
--- a/homeassistant/components/roborock/coordinator.py
+++ b/homeassistant/components/roborock/coordinator.py
@ -10,17 +10,26 @@ import logging
 from propcache.api import cached_property
 from roborock import HomeDataRoom
 from roborock.code_mappings import RoborockCategory
-from roborock.containers import DeviceData, HomeDataDevice, HomeDataProduct, NetworkInfo
+from roborock.containers import (
+    DeviceData,
+    HomeDataDevice,
+    HomeDataProduct,
+    HomeDataScene,
+    NetworkInfo,
+    UserData,
+)
 from roborock.exceptions import RoborockException
 from roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol
 from roborock.roborock_typing import DeviceProp
 from roborock.version_1_apis.roborock_local_client_v1 import RoborockLocalClientV1
 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
 from roborock.version_a01_apis import RoborockClientA01
+from roborock.web_api import RoborockApiClient

 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_CONNECTIONS
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.typing import StateType
@ -67,6 +76,8 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
        product_info: HomeDataProduct,
        cloud_api: RoborockMqttClientV1,
        home_data_rooms: list[HomeDataRoom],
+        api_client: RoborockApiClient,
+        user_data: UserData,
    ) -> None:
        """Initialize."""
        super().__init__(
@ -89,7 +100,7 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
        self.cloud_api = cloud_api
        self.device_info = DeviceInfo(
            name=self.roborock_device_info.device.name,
-            identifiers={(DOMAIN, self.roborock_device_info.device.duid)},
+            identifiers={(DOMAIN, self.duid)},
            manufacturer="Roborock",
            model=self.roborock_device_info.product.model,
            model_id=self.roborock_device_info.product.model,
@ -103,8 +114,10 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
        self.maps: dict[int, RoborockMapInfo] = {}
        self._home_data_rooms = {str(room.id): room.name for room in home_data_rooms}
        self.map_storage = RoborockMapStorage(
-            hass, self.config_entry.entry_id, slugify(self.duid)
+            hass, self.config_entry.entry_id, self.duid_slug
        )
+        self._user_data = user_data
+        self._api_client = api_client

    async def _async_setup(self) -> None:
        """Set up the coordinator."""
@ -134,7 +147,7 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
            except RoborockException:
                _LOGGER.warning(
                    "Using the cloud API for device %s. This is not recommended as it can lead to rate limiting. We recommend making your vacuum accessible by your Home Assistant instance",
-                    self.roborock_device_info.device.duid,
+                    self.duid,
                )
                await self.api.async_disconnect()
                # We use the cloud api if the local api fails to connect.
@ -193,6 +206,34 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
            for room in room_mapping or ()
        }

+    async def get_scenes(self) -> list[HomeDataScene]:
+        """Get scenes."""
+        try:
+            return await self._api_client.get_scenes(self._user_data, self.duid)
+        except RoborockException as err:
+            _LOGGER.error("Failed to get scenes %s", err)
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="command_failed",
+                translation_placeholders={
+                    "command": "get_scenes",
+                },
+            ) from err
+
+    async def execute_scene(self, scene_id: int) -> None:
+        """Execute scene."""
+        try:
+            await self._api_client.execute_scene(self._user_data, scene_id)
+        except RoborockException as err:
+            _LOGGER.error("Failed to execute scene %s %s", scene_id, err)
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="command_failed",
+                translation_placeholders={
+                    "command": "execute_scene",
+                },
+            ) from err
+
    @cached_property
    def duid(self) -> str:
        """Get the unique id of the device as specified by Roborock."""
--- a/homeassistant/components/roborock/scene.py
+++ b/homeassistant/components/roborock/scene.py
@ -0,0 +1,64 @@
+"""Support for Roborock scene."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from homeassistant.components.scene import Scene as SceneEntity
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import EntityDescription
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from . import RoborockConfigEntry
+from .coordinator import RoborockDataUpdateCoordinator
+from .entity import RoborockEntity
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: RoborockConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up scene platform."""
+    scene_lists = await asyncio.gather(
+        *[coordinator.get_scenes() for coordinator in config_entry.runtime_data.v1],
+    )
+    async_add_entities(
+        RoborockSceneEntity(
+            coordinator,
+            EntityDescription(
+                key=str(scene.id),
+                name=scene.name,
+            ),
+        )
+        for coordinator, scenes in zip(
+            config_entry.runtime_data.v1, scene_lists, strict=True
+        )
+        for scene in scenes
+    )
+
+
+class RoborockSceneEntity(RoborockEntity, SceneEntity):
+    """A class to define Roborock scene entities."""
+
+    entity_description: EntityDescription
+
+    def __init__(
+        self,
+        coordinator: RoborockDataUpdateCoordinator,
+        entity_description: EntityDescription,
+    ) -> None:
+        """Create a scene entity."""
+        super().__init__(
+            f"{entity_description.key}_{coordinator.duid_slug}",
+            coordinator.device_info,
+            coordinator.api,
+        )
+        self._scene_id = int(entity_description.key)
+        self._coordinator = coordinator
+        self.entity_description = entity_description
+
+    async def async_activate(self, **kwargs: Any) -> None:
+        """Activate the scene."""
+        await self._coordinator.execute_scene(self._scene_id)
--- a/tests/components/roborock/conftest.py
+++ b/tests/components/roborock/conftest.py
@ -30,6 +30,7 @@ from .mock_data import (
    MULTI_MAP_LIST,
    NETWORK_INFO,
    PROP,
+    SCENES,
    USER_DATA,
    USER_EMAIL,
 )
@ -67,8 +68,24 @@ class A01Mock(RoborockMqttClientA01):
        return {prot: self.protocol_responses[prot] for prot in dyad_data_protocols}


+@pytest.fixture(name="bypass_api_client_fixture")
+def bypass_api_client_fixture() -> None:
+    """Skip calls to the API client."""
+    with (
+        patch(
+            "homeassistant.components.roborock.RoborockApiClient.get_home_data_v2",
+            return_value=HOME_DATA,
+        ),
+        patch(
+            "homeassistant.components.roborock.RoborockApiClient.get_scenes",
+            return_value=SCENES,
+        ),
+    ):
+        yield
+
+
@pytest.fixture(name="bypass_api_fixture")
-def bypass_api_fixture() -> None:
+def bypass_api_fixture(bypass_api_client_fixture: Any) -> None:
    """Skip calls to the API."""
    with (
        patch("homeassistant.components.roborock.RoborockMqttClientV1.async_connect"),
@ -76,10 +93,6 @@ def bypass_api_fixture() -> None:
        patch(
            "homeassistant.components.roborock.coordinator.RoborockMqttClientV1._send_command"
        ),
-        patch(
-            "homeassistant.components.roborock.RoborockApiClient.get_home_data_v2",
-            return_value=HOME_DATA,
-        ),
        patch(
            "homeassistant.components.roborock.RoborockMqttClientV1.get_networking",
            return_value=NETWORK_INFO,
--- a/tests/components/roborock/mock_data.py
+++ b/tests/components/roborock/mock_data.py
@ -9,6 +9,7 @@ from roborock.containers import (
    Consumable,
    DnDTimer,
    HomeData,
+    HomeDataScene,
    MultiMapsList,
    NetworkInfo,
    S7Status,
@ -1150,3 +1151,19 @@ MAP_DATA = MapData(0, 0)
 MAP_DATA.image = ImageData(
    100, 10, 10, 10, 10, ImageConfig(), Image.new("RGB", (1, 1)), lambda p: p
 )
+
+
+SCENES = [
+    HomeDataScene.from_dict(
+        {
+            "name": "sc1",
+            "id": 12,
+        },
+    ),
+    HomeDataScene.from_dict(
+        {
+            "name": "sc2",
+            "id": 24,
+        },
+    ),
+]
--- a/tests/components/roborock/test_scene.py
+++ b/tests/components/roborock/test_scene.py
@ -0,0 +1,112 @@
+"""Test Roborock Scene platform."""
+
+from unittest.mock import ANY, patch
+
+import pytest
+from roborock import RoborockException
+
+from homeassistant.const import SERVICE_TURN_ON, Platform
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
+
+from tests.common import MockConfigEntry
+
+
+@pytest.fixture
+def bypass_api_client_get_scenes_fixture(bypass_api_fixture) -> None:
+    """Fixture to raise when getting scenes."""
+    with (
+        patch(
+            "homeassistant.components.roborock.RoborockApiClient.get_scenes",
+            side_effect=RoborockException(),
+        ),
+    ):
+        yield
+
+
+@pytest.mark.parametrize(
+    ("entity_id"),
+    [
+        ("scene.roborock_s7_maxv_sc1"),
+        ("scene.roborock_s7_maxv_sc2"),
+    ],
+)
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_get_scenes_failure(
+    hass: HomeAssistant,
+    bypass_api_client_get_scenes_fixture,
+    setup_entry: MockConfigEntry,
+    entity_id: str,
+) -> None:
+    """Test that if scene retrieval fails, no entity is being created."""
+    # Ensure that the entity does not exist
+    assert hass.states.get(entity_id) is None
+
+
+@pytest.fixture
+def platforms() -> list[Platform]:
+    """Fixture to set platforms used in the test."""
+    return [Platform.SCENE]
+
+
+@pytest.mark.parametrize(
+    ("entity_id", "scene_id"),
+    [
+        ("scene.roborock_s7_maxv_sc1", 12),
+        ("scene.roborock_s7_maxv_sc2", 24),
+    ],
+)
+@pytest.mark.freeze_time("2023-10-30 08:50:00")
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_execute_success(
+    hass: HomeAssistant,
+    bypass_api_fixture,
+    setup_entry: MockConfigEntry,
+    entity_id: str,
+    scene_id: int,
+) -> None:
+    """Test activating the scene entities."""
+    with patch(
+        "homeassistant.components.roborock.RoborockApiClient.execute_scene"
+    ) as mock_execute_scene:
+        await hass.services.async_call(
+            "scene",
+            SERVICE_TURN_ON,
+            blocking=True,
+            target={"entity_id": entity_id},
+        )
+    mock_execute_scene.assert_called_once_with(ANY, scene_id)
+    assert hass.states.get(entity_id).state == "2023-10-30T08:50:00+00:00"
+
+
+@pytest.mark.parametrize(
+    ("entity_id", "scene_id"),
+    [
+        ("scene.roborock_s7_maxv_sc1", 12),
+    ],
+)
+@pytest.mark.freeze_time("2023-10-30 08:50:00")
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_execute_failure(
+    hass: HomeAssistant,
+    bypass_api_fixture,
+    setup_entry: MockConfigEntry,
+    entity_id: str,
+    scene_id: int,
+) -> None:
+    """Test failure while activating the scene entity."""
+    with (
+        patch(
+            "homeassistant.components.roborock.RoborockApiClient.execute_scene",
+            side_effect=RoborockException,
+        ) as mock_execute_scene,
+        pytest.raises(HomeAssistantError, match="Error while calling execute_scene"),
+    ):
+        await hass.services.async_call(
+            "scene",
+            SERVICE_TURN_ON,
+            blocking=True,
+            target={"entity_id": entity_id},
+        )
+    mock_execute_scene.assert_called_once_with(ANY, scene_id)
+    assert hass.states.get(entity_id).state == "2023-10-30T08:50:00+00:00"