Initial implementation for tplink tapo vacuums (#131965)

Co-authored-by: Steven B. <51370195+sdb9696@users.noreply.github.com>
Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

homeassistant/components/tplink/const.py

@@ -41,6 +41,7 @@ PLATFORMS: Final = [
     Platform.SENSOR,
     Platform.SIREN,
     Platform.SWITCH,
+    Platform.VACUUM,
 ]
 
 UNIT_MAPPING = {

homeassistant/components/tplink/entity.py

@@ -59,6 +59,7 @@ DEVICETYPES_WITH_SPECIALIZED_PLATFORMS = {
     DeviceType.Dimmer,
     DeviceType.Fan,
     DeviceType.Thermostat,
+    DeviceType.Vacuum,
 }
 
 # Primary features to always include even when the device type has its own platform

homeassistant/components/tplink/vacuum.py

@@ -0,0 +1,158 @@
+"""Support for TPLink vacuum."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, cast
+
+from kasa import Device, Feature, Module
+from kasa.smart.modules.clean import Clean, Status
+
+from homeassistant.components.vacuum import (
+    DOMAIN as VACUUM_DOMAIN,
+    StateVacuumEntity,
+    StateVacuumEntityDescription,
+    VacuumActivity,
+    VacuumEntityFeature,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from . import TPLinkConfigEntry
+from .coordinator import TPLinkDataUpdateCoordinator
+from .entity import (
+    CoordinatedTPLinkModuleEntity,
+    TPLinkModuleEntityDescription,
+    async_refresh_after,
+)
+
+# Coordinator is used to centralize the data updates
+# For actions the integration handles locking of concurrent device request
+PARALLEL_UPDATES = 0
+
+# Upstream state to VacuumActivity
+STATUS_TO_ACTIVITY = {
+    Status.Idle: VacuumActivity.IDLE,
+    Status.Cleaning: VacuumActivity.CLEANING,
+    Status.GoingHome: VacuumActivity.RETURNING,
+    Status.Charging: VacuumActivity.DOCKED,
+    Status.Charged: VacuumActivity.DOCKED,
+    Status.Undocked: VacuumActivity.IDLE,
+    Status.Paused: VacuumActivity.PAUSED,
+    Status.Error: VacuumActivity.ERROR,
+}
+
+
+@dataclass(frozen=True, kw_only=True)
+class TPLinkVacuumEntityDescription(
+    StateVacuumEntityDescription, TPLinkModuleEntityDescription
+):
+    """Base class for vacuum entity description."""
+
+
+VACUUM_DESCRIPTIONS: tuple[TPLinkVacuumEntityDescription, ...] = (
+    TPLinkVacuumEntityDescription(
+        key="vacuum", exists_fn=lambda dev, _: Module.Clean in dev.modules
+    ),
+)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: TPLinkConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up vacuum entities."""
+    data = config_entry.runtime_data
+    parent_coordinator = data.parent_coordinator
+    device = parent_coordinator.device
+
+    known_child_device_ids: set[str] = set()
+    first_check = True
+
+    def _check_device() -> None:
+        entities = CoordinatedTPLinkModuleEntity.entities_for_device_and_its_children(
+            hass=hass,
+            device=device,
+            coordinator=parent_coordinator,
+            entity_class=TPLinkVacuumEntity,
+            descriptions=VACUUM_DESCRIPTIONS,
+            platform_domain=VACUUM_DOMAIN,
+            known_child_device_ids=known_child_device_ids,
+            first_check=first_check,
+        )
+        async_add_entities(entities)
+
+    _check_device()
+    first_check = False
+    config_entry.async_on_unload(parent_coordinator.async_add_listener(_check_device))
+
+
+class TPLinkVacuumEntity(CoordinatedTPLinkModuleEntity, StateVacuumEntity):
+    """Representation of a tplink vacuum."""
+
+    _attr_supported_features = (
+        VacuumEntityFeature.STATE
+        | VacuumEntityFeature.BATTERY
+        | VacuumEntityFeature.START
+        | VacuumEntityFeature.PAUSE
+        | VacuumEntityFeature.RETURN_HOME
+        | VacuumEntityFeature.FAN_SPEED
+    )
+
+    entity_description: TPLinkVacuumEntityDescription
+
+    def __init__(
+        self,
+        device: Device,
+        coordinator: TPLinkDataUpdateCoordinator,
+        description: TPLinkVacuumEntityDescription,
+        *,
+        parent: Device,
+    ) -> None:
+        """Initialize the vacuum entity."""
+        super().__init__(device, coordinator, description, parent=parent)
+        self._vacuum_module: Clean = device.modules[Module.Clean]
+        if speaker := device.modules.get(Module.Speaker):
+            self._speaker_module = speaker
+            self._attr_supported_features |= VacuumEntityFeature.LOCATE
+
+        # Needs to be initialized empty, as vacuumentity's capability_attributes accesses it
+        self._attr_fan_speed_list: list[str] = []
+
+    @async_refresh_after
+    async def async_start(self) -> None:
+        """Start cleaning."""
+        await self._vacuum_module.start()
+
+    @async_refresh_after
+    async def async_pause(self) -> None:
+        """Pause cleaning."""
+        await self._vacuum_module.pause()
+
+    @async_refresh_after
+    async def async_return_to_base(self, **kwargs: Any) -> None:
+        """Return home."""
+        await self._vacuum_module.return_home()
+
+    @async_refresh_after
+    async def async_set_fan_speed(self, fan_speed: str, **kwargs: Any) -> None:
+        """Set fan speed."""
+        await self._vacuum_module.set_fan_speed_preset(fan_speed)
+
+    async def async_locate(self, **kwargs: Any) -> None:
+        """Locate the device."""
+        await self._speaker_module.locate()
+
+    @property
+    def battery_level(self) -> int | None:
+        """Return battery level."""
+        return self._vacuum_module.battery
+
+    def _async_update_attrs(self) -> bool:
+        """Update the entity's attributes."""
+        self._attr_activity = STATUS_TO_ACTIVITY.get(self._vacuum_module.status)
+        fanspeeds = cast(Feature, self._vacuum_module.get_feature("fan_speed_preset"))
+        self._attr_fan_speed_list = cast(list[str], fanspeeds.choices)
+        self._attr_fan_speed = self._vacuum_module.fan_speed_preset
+        return True

tests/components/tplink/__init__.py

@@ -16,7 +16,9 @@ from kasa import (
     ThermostatState,
 )
 from kasa.interfaces import Fan, Light, LightEffect, LightState, Thermostat
+from kasa.smart.modules import Speaker
 from kasa.smart.modules.alarm import Alarm
+from kasa.smart.modules.clean import Clean, ErrorCode, Status
 from kasa.smartcam.modules.camera import LOCAL_STREAMING_PORT, Camera
 from syrupy import SnapshotAssertion
 
@@ -195,16 +197,33 @@ def _mocked_device(
         )
     device.features = device_features
 
-    # Add modules after features so modules can add required features
+    # Add modules after features so modules can add any required features
     if modules:
         device.modules = {
             module_name: MODULE_TO_MOCK_GEN[module_name](device)
             for module_name in modules
         }
 
+    # module features are accessed from a module via get_feature which is
+    # keyed on the module attribute name. Usually this is the same as the
+    # feature.id but not always so accept overrides.
+    module_features = {
+        mod_key if (mod_key := v.expected_module_key) else k: v
+        for k, v in device_features.items()
+    }
     for mod in device.modules.values():
-        mod.get_feature.side_effect = device_features.get
-        mod.has_feature.side_effect = lambda id: id in device_features
+        # Some tests remove the feature from device_features to test missing
+        # features, so check the key is still present there.
+        mod.get_feature.side_effect = (
+            lambda mod_id: mod_feat
+            if (mod_feat := module_features.get(mod_id))
+            and mod_feat.id in device_features
+            else None
+        )
+        mod.has_feature.side_effect = (
+            lambda mod_id: (mod_feat := module_features.get(mod_id))
+            and mod_feat.id in device_features
+        )
 
     device.parent = None
     device.children = []
@@ -243,6 +262,7 @@ def _mocked_feature(
     unit=None,
     minimum_value=None,
     maximum_value=None,
+    expected_module_key=None,
 ) -> Feature:
     """Get a mocked feature.
 
@@ -284,6 +304,16 @@ def _mocked_feature(
     # select
     feature.choices = choices or fixture.get("choices")
 
+    # module features are accessed from a module via get_feature which is
+    # keyed on the module attribute name. Usually this is the same as the
+    # feature.id but not always. module_key indicates the key of the feature
+    # in the module.
+    feature.expected_module_key = (
+        mod_key
+        if (mod_key := fixture.get("expected_module_key", expected_module_key))
+        else None
+    )
+
     return feature
 
 
@@ -400,6 +430,43 @@ def _mocked_thermostat_module(device):
     return therm
 
 
+def _mocked_clean_module(device):
+    clean = MagicMock(auto_spec=Clean, name="Mocked clean")
+
+    # methods
+    clean.start = AsyncMock()
+    clean.pause = AsyncMock()
+    clean.resume = AsyncMock()
+    clean.return_home = AsyncMock()
+    clean.set_fan_speed_preset = AsyncMock()
+
+    # properties
+    clean.fan_speed_preset = "Max"
+    clean.error = ErrorCode.Ok
+    clean.battery = 100
+    clean.status = Status.Charged
+
+    # Need to manually create the fan speed preset feature,
+    # as we are going to read its choices through it
+    device.features["vacuum_fan_speed"] = _mocked_feature(
+        "vacuum_fan_speed",
+        type_=Feature.Type.Choice,
+        category=Feature.Category.Config,
+        choices=["Quiet", "Max"],
+        value="Max",
+        expected_module_key="fan_speed_preset",
+    )
+
+    return clean
+
+
+def _mocked_speaker_module(device):
+    speaker = MagicMock(auto_spec=Speaker, name="Mocked speaker")
+    speaker.locate = AsyncMock()
+
+    return speaker
+
+
 def _mocked_strip_children(features=None, alias=None) -> list[Device]:
     plug0 = _mocked_device(
         alias="Plug0" if alias is None else alias,
@@ -469,6 +536,8 @@ MODULE_TO_MOCK_GEN = {
     Module.Alarm: _mocked_alarm_module,
     Module.Camera: _mocked_camera_module,
     Module.Thermostat: _mocked_thermostat_module,
+    Module.Clean: _mocked_clean_module,
+    Module.Speaker: _mocked_speaker_module,
 }
 
 

tests/components/tplink/snapshots/test_vacuum.ambr

@@ -0,0 +1,96 @@
+# serializer version: 1
+# name: test_states[my_vacuum-entry]
+  DeviceRegistryEntrySnapshot({
+    'area_id': None,
+    'config_entries': <ANY>,
+    'configuration_url': None,
+    'connections': set({
+      tuple(
+        'mac',
+        'aa:bb:cc:dd:ee:ff',
+      ),
+    }),
+    'disabled_by': None,
+    'entry_type': None,
+    'hw_version': '1.0.0',
+    'id': <ANY>,
+    'identifiers': set({
+      tuple(
+        'tplink',
+        '123456789ABCDEFGH',
+      ),
+    }),
+    'is_new': False,
+    'labels': set({
+    }),
+    'manufacturer': 'TP-Link',
+    'model': 'HS100',
+    'model_id': None,
+    'name': 'my_vacuum',
+    'name_by_user': None,
+    'primary_config_entry': <ANY>,
+    'serial_number': None,
+    'suggested_area': None,
+    'sw_version': '1.0.0',
+    'via_device_id': None,
+  })
+# ---
+# name: test_states[vacuum.my_vacuum-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': dict({
+      'fan_speed_list': list([
+        'Quiet',
+        'Max',
+      ]),
+    }),
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'vacuum',
+    'entity_category': None,
+    'entity_id': 'vacuum.my_vacuum',
+    'has_entity_name': True,
+    'hidden_by': None,
+    'icon': None,
+    'id': <ANY>,
+    'labels': set({
+    }),
+    'name': None,
+    'options': dict({
+    }),
+    'original_device_class': None,
+    'original_icon': None,
+    'original_name': None,
+    'platform': 'tplink',
+    'previous_unique_id': None,
+    'supported_features': <VacuumEntityFeature: 12916>,
+    'translation_key': None,
+    'unique_id': '123456789ABCDEFGH-vacuum',
+    'unit_of_measurement': None,
+  })
+# ---
+# name: test_states[vacuum.my_vacuum-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'battery_icon': 'mdi:battery-charging-100',
+      'battery_level': 100,
+      'fan_speed': 'Max',
+      'fan_speed_list': list([
+        'Quiet',
+        'Max',
+      ]),
+      'friendly_name': 'my_vacuum',
+      'supported_features': <VacuumEntityFeature: 12916>,
+    }),
+    'context': <ANY>,
+    'entity_id': 'vacuum.my_vacuum',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'docked',
+  })
+# ---

tests/components/tplink/test_vacuum.py

@@ -0,0 +1,125 @@
+"""Tests for vacuum platform."""
+
+from __future__ import annotations
+
+from kasa import Device, Module
+import pytest
+from syrupy.assertion import SnapshotAssertion
+
+from homeassistant.components.vacuum import (
+    ATTR_BATTERY_LEVEL,
+    ATTR_FAN_SPEED,
+    DOMAIN as VACUUM_DOMAIN,
+    SERVICE_LOCATE,
+    SERVICE_PAUSE,
+    SERVICE_RETURN_TO_BASE,
+    SERVICE_SET_FAN_SPEED,
+    SERVICE_START,
+    VacuumActivity,
+)
+from homeassistant.const import ATTR_ENTITY_ID, Platform
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr, entity_registry as er
+
+from . import DEVICE_ID, _mocked_device, setup_platform_for_device, snapshot_platform
+
+from tests.common import MockConfigEntry
+
+ENTITY_ID = "vacuum.my_vacuum"
+
+
+@pytest.fixture
+async def mocked_vacuum(hass: HomeAssistant) -> Device:
+    """Return mocked tplink vacuum."""
+
+    return _mocked_device(modules=[Module.Clean, Module.Speaker], alias="my_vacuum")
+
+
+async def test_vacuum(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    entity_registry: er.EntityRegistry,
+    device_registry: dr.DeviceRegistry,
+    mocked_vacuum: Device,
+) -> None:
+    """Test initialization."""
+    await setup_platform_for_device(
+        hass, mock_config_entry, Platform.VACUUM, mocked_vacuum
+    )
+
+    device_entries = dr.async_entries_for_config_entry(
+        device_registry, mock_config_entry.entry_id
+    )
+    assert device_entries
+
+    entity = entity_registry.async_get(ENTITY_ID)
+    assert entity
+    assert entity.unique_id == f"{DEVICE_ID}-vacuum"
+
+    state = hass.states.get(ENTITY_ID)
+    assert state.state == VacuumActivity.DOCKED
+
+    assert state.attributes[ATTR_FAN_SPEED] == "Max"
+    assert state.attributes[ATTR_BATTERY_LEVEL] == 100
+
+
+async def test_states(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    entity_registry: er.EntityRegistry,
+    device_registry: dr.DeviceRegistry,
+    snapshot: SnapshotAssertion,
+    mocked_vacuum: Device,
+) -> None:
+    """Test vacuum states."""
+    await setup_platform_for_device(
+        hass, mock_config_entry, Platform.VACUUM, mocked_vacuum
+    )
+    await snapshot_platform(
+        hass, entity_registry, device_registry, snapshot, mock_config_entry.entry_id
+    )
+
+
+@pytest.mark.parametrize(
+    ("service_call", "module_name", "method", "params"),
+    [
+        (SERVICE_START, Module.Clean, "start", {}),
+        (SERVICE_PAUSE, Module.Clean, "pause", {}),
+        (SERVICE_RETURN_TO_BASE, Module.Clean, "return_home", {}),
+        (
+            SERVICE_SET_FAN_SPEED,
+            Module.Clean,
+            "set_fan_speed_preset",
+            {ATTR_FAN_SPEED: "Quiet"},
+        ),
+        (SERVICE_LOCATE, Module.Speaker, "locate", {}),
+    ],
+)
+async def test_vacuum_module(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mocked_vacuum: Device,
+    service_call: str,
+    module_name: str,
+    method: str,
+    params: dict,
+) -> None:
+    """Test that all vacuum commands work correctly."""
+    vacuum = mocked_vacuum
+    module = vacuum.modules[module_name]
+
+    await setup_platform_for_device(hass, mock_config_entry, Platform.VACUUM, vacuum)
+
+    mock_method = getattr(module, method)
+
+    service_data = {ATTR_ENTITY_ID: ENTITY_ID}
+    service_data |= params
+
+    await hass.services.async_call(
+        VACUUM_DOMAIN, service_call, service_data, blocking=True
+    )
+
+    # Is this required when using blocking=True?
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    mock_method.assert_called()