diff --git a/homeassistant/components/tplink/const.py b/homeassistant/components/tplink/const.py index 61c1bf1cb9b..ad17aadeb5b 100644 --- a/homeassistant/components/tplink/const.py +++ b/homeassistant/components/tplink/const.py @@ -41,6 +41,7 @@ PLATFORMS: Final = [ Platform.SENSOR, Platform.SIREN, Platform.SWITCH, + Platform.VACUUM, ] UNIT_MAPPING = { diff --git a/homeassistant/components/tplink/entity.py b/homeassistant/components/tplink/entity.py index edef8bd83a0..6c21ab63285 100644 --- a/homeassistant/components/tplink/entity.py +++ b/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 diff --git a/homeassistant/components/tplink/vacuum.py b/homeassistant/components/tplink/vacuum.py new file mode 100644 index 00000000000..666584f4980 --- /dev/null +++ b/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 diff --git a/tests/components/tplink/__init__.py b/tests/components/tplink/__init__.py index 008d25a3dcb..664fb96fe71 100644 --- a/tests/components/tplink/__init__.py +++ b/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, } diff --git a/tests/components/tplink/snapshots/test_vacuum.ambr b/tests/components/tplink/snapshots/test_vacuum.ambr new file mode 100644 index 00000000000..a28a7d80ab4 --- /dev/null +++ b/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': , + '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': , + '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': , + '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': , + 'device_class': None, + 'device_id': , + 'disabled_by': None, + 'domain': 'vacuum', + 'entity_category': None, + 'entity_id': 'vacuum.my_vacuum', + 'has_entity_name': True, + 'hidden_by': None, + 'icon': None, + 'id': , + 'labels': set({ + }), + 'name': None, + 'options': dict({ + }), + 'original_device_class': None, + 'original_icon': None, + 'original_name': None, + 'platform': 'tplink', + 'previous_unique_id': None, + 'supported_features': , + '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': , + }), + 'context': , + 'entity_id': 'vacuum.my_vacuum', + 'last_changed': , + 'last_reported': , + 'last_updated': , + 'state': 'docked', + }) +# --- diff --git a/tests/components/tplink/test_vacuum.py b/tests/components/tplink/test_vacuum.py new file mode 100644 index 00000000000..aac7c4f7fc8 --- /dev/null +++ b/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()