diff --git a/homeassistant/components/husqvarna_automower/__init__.py b/homeassistant/components/husqvarna_automower/__init__.py index da7965250cd..a08256fb0b5 100644 --- a/homeassistant/components/husqvarna_automower/__init__.py +++ b/homeassistant/components/husqvarna_automower/__init__.py @@ -9,16 +9,10 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady -from homeassistant.helpers import ( - aiohttp_client, - config_entry_oauth2_flow, - device_registry as dr, - entity_registry as er, -) +from homeassistant.helpers import aiohttp_client, config_entry_oauth2_flow from homeassistant.util import dt as dt_util from . import api -from .const import DOMAIN from .coordinator import AutomowerDataUpdateCoordinator _LOGGER = logging.getLogger(__name__) @@ -69,8 +63,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: AutomowerConfigEntry) -> coordinator = AutomowerDataUpdateCoordinator(hass, automower_api) await coordinator.async_config_entry_first_refresh() - available_devices = list(coordinator.data) - cleanup_removed_devices(hass, coordinator.config_entry, available_devices) entry.runtime_data = coordinator entry.async_create_background_task( @@ -86,36 +78,3 @@ async def async_setup_entry(hass: HomeAssistant, entry: AutomowerConfigEntry) -> async def async_unload_entry(hass: HomeAssistant, entry: AutomowerConfigEntry) -> bool: """Handle unload of an entry.""" return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - - -def cleanup_removed_devices( - hass: HomeAssistant, - config_entry: AutomowerConfigEntry, - available_devices: list[str], -) -> None: - """Cleanup entity and device registry from removed devices.""" - device_reg = dr.async_get(hass) - identifiers = {(DOMAIN, mower_id) for mower_id in available_devices} - for device in dr.async_entries_for_config_entry(device_reg, config_entry.entry_id): - if not set(device.identifiers) & identifiers: - _LOGGER.debug("Removing obsolete device entry %s", device.name) - device_reg.async_update_device( - device.id, remove_config_entry_id=config_entry.entry_id - ) - - -def remove_work_area_entities( - hass: HomeAssistant, - config_entry: AutomowerConfigEntry, - removed_work_areas: set[int], - mower_id: str, -) -> None: - """Remove all unused work area entities for the specified mower.""" - entity_reg = er.async_get(hass) - for entity_entry in er.async_entries_for_config_entry( - entity_reg, config_entry.entry_id - ): - for work_area_id in removed_work_areas: - if entity_entry.unique_id.startswith(f"{mower_id}_{work_area_id}_"): - _LOGGER.info("Deleting: %s", entity_entry.entity_id) - entity_reg.async_remove(entity_entry.entity_id) diff --git a/homeassistant/components/husqvarna_automower/binary_sensor.py b/homeassistant/components/husqvarna_automower/binary_sensor.py index 3c23da76797..907d34e812a 100644 --- a/homeassistant/components/husqvarna_automower/binary_sensor.py +++ b/homeassistant/components/husqvarna_automower/binary_sensor.py @@ -75,11 +75,16 @@ async def async_setup_entry( ) -> None: """Set up binary sensor platform.""" coordinator = entry.runtime_data - async_add_entities( - AutomowerBinarySensorEntity(mower_id, coordinator, description) - for mower_id in coordinator.data - for description in MOWER_BINARY_SENSOR_TYPES - ) + + def _async_add_new_devices(mower_ids: set[str]) -> None: + async_add_entities( + AutomowerBinarySensorEntity(mower_id, coordinator, description) + for mower_id in mower_ids + for description in MOWER_BINARY_SENSOR_TYPES + ) + + coordinator.new_devices_callbacks.append(_async_add_new_devices) + _async_add_new_devices(set(coordinator.data)) class AutomowerBinarySensorEntity(AutomowerBaseEntity, BinarySensorEntity): diff --git a/homeassistant/components/husqvarna_automower/button.py b/homeassistant/components/husqvarna_automower/button.py index ce303325496..7e6e581cdf1 100644 --- a/homeassistant/components/husqvarna_automower/button.py +++ b/homeassistant/components/husqvarna_automower/button.py @@ -58,12 +58,17 @@ async def async_setup_entry( ) -> None: """Set up button platform.""" coordinator = entry.runtime_data - async_add_entities( - AutomowerButtonEntity(mower_id, coordinator, description) - for mower_id in coordinator.data - for description in MOWER_BUTTON_TYPES - if description.exists_fn(coordinator.data[mower_id]) - ) + + def _async_add_new_devices(mower_ids: set[str]) -> None: + async_add_entities( + AutomowerButtonEntity(mower_id, coordinator, description) + for mower_id in mower_ids + for description in MOWER_BUTTON_TYPES + if description.exists_fn(coordinator.data[mower_id]) + ) + + coordinator.new_devices_callbacks.append(_async_add_new_devices) + _async_add_new_devices(set(coordinator.data)) class AutomowerButtonEntity(AutomowerAvailableEntity, ButtonEntity): diff --git a/homeassistant/components/husqvarna_automower/calendar.py b/homeassistant/components/husqvarna_automower/calendar.py index f3e82fde5d4..9e2ea037afb 100644 --- a/homeassistant/components/husqvarna_automower/calendar.py +++ b/homeassistant/components/husqvarna_automower/calendar.py @@ -26,9 +26,14 @@ async def async_setup_entry( ) -> None: """Set up lawn mower platform.""" coordinator = entry.runtime_data - async_add_entities( - AutomowerCalendarEntity(mower_id, coordinator) for mower_id in coordinator.data - ) + + def _async_add_new_devices(mower_ids: set[str]) -> None: + async_add_entities( + AutomowerCalendarEntity(mower_id, coordinator) for mower_id in mower_ids + ) + + coordinator.new_devices_callbacks.append(_async_add_new_devices) + _async_add_new_devices(set(coordinator.data)) class AutomowerCalendarEntity(AutomowerBaseEntity, CalendarEntity): diff --git a/homeassistant/components/husqvarna_automower/coordinator.py b/homeassistant/components/husqvarna_automower/coordinator.py index 57be02e7066..2921b5ca68e 100644 --- a/homeassistant/components/husqvarna_automower/coordinator.py +++ b/homeassistant/components/husqvarna_automower/coordinator.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from collections.abc import Callable from datetime import timedelta import logging from typing import TYPE_CHECKING @@ -18,6 +19,7 @@ from aioautomower.session import AutomowerSession from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import ConfigEntryAuthFailed +from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from .const import DOMAIN @@ -47,6 +49,12 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[dict[str, MowerAttrib self.api = api self.ws_connected: bool = False self.reconnect_time = DEFAULT_RECONNECT_TIME + self.new_devices_callbacks: list[Callable[[set[str]], None]] = [] + self.new_zones_callbacks: list[Callable[[str, set[str]], None]] = [] + self.new_areas_callbacks: list[Callable[[str, set[int]], None]] = [] + self._devices_last_update: set[str] = set() + self._zones_last_update: dict[str, set[str]] = {} + self._areas_last_update: dict[str, set[int]] = {} async def _async_update_data(self) -> dict[str, MowerAttributes]: """Subscribe for websocket and poll data from the API.""" @@ -55,12 +63,21 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[dict[str, MowerAttrib self.api.register_data_callback(self.callback) self.ws_connected = True try: - return await self.api.get_status() + data = await self.api.get_status() except ApiException as err: raise UpdateFailed(err) from err except AuthException as err: raise ConfigEntryAuthFailed(err) from err + self._async_add_remove_devices(data) + for mower_id in data: + if data[mower_id].capabilities.stay_out_zones: + self._async_add_remove_stay_out_zones(data) + for mower_id in data: + if data[mower_id].capabilities.work_areas: + self._async_add_remove_work_areas(data) + return data + @callback def callback(self, ws_data: dict[str, MowerAttributes]) -> None: """Process websocket callbacks and write them to the DataUpdateCoordinator.""" @@ -96,3 +113,136 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[dict[str, MowerAttrib self.client_listen(hass, entry, automower_client), "reconnect_task", ) + + def _async_add_remove_devices(self, data: dict[str, MowerAttributes]) -> None: + """Add new device, remove non-existing device.""" + current_devices = set(data) + + # Skip update if no changes + if current_devices == self._devices_last_update: + return + + # Process removed devices + removed_devices = self._devices_last_update - current_devices + if removed_devices: + _LOGGER.debug("Removed devices: %s", ", ".join(map(str, removed_devices))) + self._remove_device(removed_devices) + + # Process new device + new_devices = current_devices - self._devices_last_update + if new_devices: + _LOGGER.debug("New devices found: %s", ", ".join(map(str, new_devices))) + self._add_new_devices(new_devices) + + # Update device state + self._devices_last_update = current_devices + + def _remove_device(self, removed_devices: set[str]) -> None: + """Remove device from the registry.""" + device_registry = dr.async_get(self.hass) + for mower_id in removed_devices: + if device := device_registry.async_get_device( + identifiers={(DOMAIN, str(mower_id))} + ): + device_registry.async_update_device( + device_id=device.id, + remove_config_entry_id=self.config_entry.entry_id, + ) + + def _add_new_devices(self, new_devices: set[str]) -> None: + """Add new device and trigger callbacks.""" + for mower_callback in self.new_devices_callbacks: + mower_callback(new_devices) + + def _async_add_remove_stay_out_zones( + self, data: dict[str, MowerAttributes] + ) -> None: + """Add new stay-out zones, remove non-existing stay-out zones.""" + current_zones = { + mower_id: set(mower_data.stay_out_zones.zones) + for mower_id, mower_data in data.items() + if mower_data.capabilities.stay_out_zones + and mower_data.stay_out_zones is not None + } + + if not self._zones_last_update: + self._zones_last_update = current_zones + return + + if current_zones == self._zones_last_update: + return + + self._zones_last_update = self._update_stay_out_zones(current_zones) + + def _update_stay_out_zones( + self, current_zones: dict[str, set[str]] + ) -> dict[str, set[str]]: + """Update stay-out zones by adding and removing as needed.""" + new_zones = { + mower_id: zones - self._zones_last_update.get(mower_id, set()) + for mower_id, zones in current_zones.items() + } + removed_zones = { + mower_id: self._zones_last_update.get(mower_id, set()) - zones + for mower_id, zones in current_zones.items() + } + + for mower_id, zones in new_zones.items(): + for zone_callback in self.new_zones_callbacks: + zone_callback(mower_id, set(zones)) + + entity_registry = er.async_get(self.hass) + for mower_id, zones in removed_zones.items(): + for entity_entry in er.async_entries_for_config_entry( + entity_registry, self.config_entry.entry_id + ): + for zone in zones: + if entity_entry.unique_id.startswith(f"{mower_id}_{zone}"): + entity_registry.async_remove(entity_entry.entity_id) + + return current_zones + + def _async_add_remove_work_areas(self, data: dict[str, MowerAttributes]) -> None: + """Add new work areas, remove non-existing work areas.""" + current_areas = { + mower_id: set(mower_data.work_areas) + for mower_id, mower_data in data.items() + if mower_data.capabilities.work_areas and mower_data.work_areas is not None + } + + if not self._areas_last_update: + self._areas_last_update = current_areas + return + + if current_areas == self._areas_last_update: + return + + self._areas_last_update = self._update_work_areas(current_areas) + + def _update_work_areas( + self, current_areas: dict[str, set[int]] + ) -> dict[str, set[int]]: + """Update work areas by adding and removing as needed.""" + new_areas = { + mower_id: areas - self._areas_last_update.get(mower_id, set()) + for mower_id, areas in current_areas.items() + } + removed_areas = { + mower_id: self._areas_last_update.get(mower_id, set()) - areas + for mower_id, areas in current_areas.items() + } + + for mower_id, areas in new_areas.items(): + for area_callback in self.new_areas_callbacks: + area_callback(mower_id, set(areas)) + + entity_registry = er.async_get(self.hass) + for mower_id, areas in removed_areas.items(): + for entity_entry in er.async_entries_for_config_entry( + entity_registry, self.config_entry.entry_id + ): + for area in areas: + if entity_entry.unique_id.startswith(f"{mower_id}_{area}_"): + entity_registry.async_remove(entity_entry.entity_id) + + return current_areas diff --git a/homeassistant/components/husqvarna_automower/device_tracker.py b/homeassistant/components/husqvarna_automower/device_tracker.py index 520eaceb1d0..2fd59b63014 100644 --- a/homeassistant/components/husqvarna_automower/device_tracker.py +++ b/homeassistant/components/husqvarna_automower/device_tracker.py @@ -19,11 +19,16 @@ async def async_setup_entry( ) -> None: """Set up device tracker platform.""" coordinator = entry.runtime_data - async_add_entities( - AutomowerDeviceTrackerEntity(mower_id, coordinator) - for mower_id in coordinator.data - if coordinator.data[mower_id].capabilities.position - ) + + def _async_add_new_devices(mower_ids: set[str]) -> None: + async_add_entities( + AutomowerDeviceTrackerEntity(mower_id, coordinator) + for mower_id in mower_ids + if coordinator.data[mower_id].capabilities.position + ) + + coordinator.new_devices_callbacks.append(_async_add_new_devices) + _async_add_new_devices(set(coordinator.data)) class AutomowerDeviceTrackerEntity(AutomowerBaseEntity, TrackerEntity): diff --git a/homeassistant/components/husqvarna_automower/lawn_mower.py b/homeassistant/components/husqvarna_automower/lawn_mower.py index 9b3ce7dab1a..dd75a8b9bc4 100644 --- a/homeassistant/components/husqvarna_automower/lawn_mower.py +++ b/homeassistant/components/husqvarna_automower/lawn_mower.py @@ -53,10 +53,15 @@ async def async_setup_entry( ) -> None: """Set up lawn mower platform.""" coordinator = entry.runtime_data - async_add_entities( - AutomowerLawnMowerEntity(mower_id, coordinator) for mower_id in coordinator.data - ) + def _async_add_new_devices(mower_ids: set[str]) -> None: + async_add_entities( + [AutomowerLawnMowerEntity(mower_id, coordinator) for mower_id in mower_ids] + ) + + _async_add_new_devices(set(coordinator.data)) + + coordinator.new_devices_callbacks.append(_async_add_new_devices) platform = entity_platform.async_get_current_platform() platform.async_register_entity_service( "override_schedule", diff --git a/homeassistant/components/husqvarna_automower/number.py b/homeassistant/components/husqvarna_automower/number.py index e69b52fab93..d3666494646 100644 --- a/homeassistant/components/husqvarna_automower/number.py +++ b/homeassistant/components/husqvarna_automower/number.py @@ -13,7 +13,7 @@ from homeassistant.const import PERCENTAGE, EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import AutomowerConfigEntry, remove_work_area_entities +from . import AutomowerConfigEntry from .coordinator import AutomowerDataUpdateCoordinator from .entity import ( AutomowerControlEntity, @@ -111,44 +111,47 @@ async def async_setup_entry( ) -> None: """Set up number platform.""" coordinator = entry.runtime_data - current_work_areas: dict[str, set[int]] = {} - - async_add_entities( - AutomowerNumberEntity(mower_id, coordinator, description) - for mower_id in coordinator.data - for description in MOWER_NUMBER_TYPES - if description.exists_fn(coordinator.data[mower_id]) - ) - - def _async_work_area_listener() -> None: - """Listen for new work areas and add/remove entities as needed.""" - for mower_id in coordinator.data: - if ( - coordinator.data[mower_id].capabilities.work_areas - and (_work_areas := coordinator.data[mower_id].work_areas) is not None - ): - received_work_areas = set(_work_areas.keys()) - current_work_area_set = current_work_areas.setdefault(mower_id, set()) - - new_work_areas = received_work_areas - current_work_area_set - removed_work_areas = current_work_area_set - received_work_areas - - if new_work_areas: - current_work_area_set.update(new_work_areas) - async_add_entities( - WorkAreaNumberEntity( - mower_id, coordinator, description, work_area_id - ) - for description in WORK_AREA_NUMBER_TYPES - for work_area_id in new_work_areas + entities: list[NumberEntity] = [] + for mower_id in coordinator.data: + if coordinator.data[mower_id].capabilities.work_areas: + _work_areas = coordinator.data[mower_id].work_areas + if _work_areas is not None: + entities.extend( + WorkAreaNumberEntity( + mower_id, coordinator, description, work_area_id ) + for description in WORK_AREA_NUMBER_TYPES + for work_area_id in _work_areas + ) + entities.extend( + AutomowerNumberEntity(mower_id, coordinator, description) + for description in MOWER_NUMBER_TYPES + if description.exists_fn(coordinator.data[mower_id]) + ) + async_add_entities(entities) - if removed_work_areas: - remove_work_area_entities(hass, entry, removed_work_areas, mower_id) - current_work_area_set.difference_update(removed_work_areas) + def _async_add_new_work_areas(mower_id: str, work_area_ids: set[int]) -> None: + async_add_entities( + WorkAreaNumberEntity(mower_id, coordinator, description, work_area_id) + for description in WORK_AREA_NUMBER_TYPES + for work_area_id in work_area_ids + ) - coordinator.async_add_listener(_async_work_area_listener) - _async_work_area_listener() + def _async_add_new_devices(mower_ids: set[str]) -> None: + async_add_entities( + AutomowerNumberEntity(mower_id, coordinator, description) + for description in MOWER_NUMBER_TYPES + for mower_id in mower_ids + if description.exists_fn(coordinator.data[mower_id]) + ) + for mower_id in mower_ids: + mower_data = coordinator.data[mower_id] + if mower_data.capabilities.work_areas and mower_data.work_areas is not None: + work_area_ids = set(mower_data.work_areas.keys()) + _async_add_new_work_areas(mower_id, work_area_ids) + + coordinator.new_areas_callbacks.append(_async_add_new_work_areas) + coordinator.new_devices_callbacks.append(_async_add_new_devices) class AutomowerNumberEntity(AutomowerControlEntity, NumberEntity): diff --git a/homeassistant/components/husqvarna_automower/quality_scale.yaml b/homeassistant/components/husqvarna_automower/quality_scale.yaml index 2287ccb4d4f..2fa41c02a4c 100644 --- a/homeassistant/components/husqvarna_automower/quality_scale.yaml +++ b/homeassistant/components/husqvarna_automower/quality_scale.yaml @@ -57,9 +57,7 @@ rules: docs-supported-functions: done docs-troubleshooting: done docs-use-cases: todo - dynamic-devices: - status: todo - comment: Add devices dynamically + dynamic-devices: done entity-category: done entity-device-class: done entity-disabled-by-default: done @@ -70,9 +68,7 @@ rules: status: exempt comment: no configuration possible repair-issues: done - stale-devices: - status: todo - comment: We only remove devices on reload + stale-devices: done # Platinum async-dependency: done diff --git a/homeassistant/components/husqvarna_automower/select.py b/homeassistant/components/husqvarna_automower/select.py index 65960e897e4..03b1ac02587 100644 --- a/homeassistant/components/husqvarna_automower/select.py +++ b/homeassistant/components/husqvarna_automower/select.py @@ -33,11 +33,17 @@ async def async_setup_entry( ) -> None: """Set up select platform.""" coordinator = entry.runtime_data - async_add_entities( - AutomowerSelectEntity(mower_id, coordinator) - for mower_id in coordinator.data - if coordinator.data[mower_id].capabilities.headlights - ) + + def _async_add_new_devices(mower_ids: set[str]) -> None: + async_add_entities( + AutomowerSelectEntity(mower_id, coordinator) + for mower_id in mower_ids + if coordinator.data[mower_id].capabilities.headlights + ) + + _async_add_new_devices(set(coordinator.data)) + + coordinator.new_devices_callbacks.append(_async_add_new_devices) class AutomowerSelectEntity(AutomowerControlEntity, SelectEntity): diff --git a/homeassistant/components/husqvarna_automower/sensor.py b/homeassistant/components/husqvarna_automower/sensor.py index fb8603623e4..a2f4b5f4bab 100644 --- a/homeassistant/components/husqvarna_automower/sensor.py +++ b/homeassistant/components/husqvarna_automower/sensor.py @@ -434,44 +434,56 @@ async def async_setup_entry( ) -> None: """Set up sensor platform.""" coordinator = entry.runtime_data - current_work_areas: dict[str, set[int]] = {} - - async_add_entities( - AutomowerSensorEntity(mower_id, coordinator, description) - for mower_id, data in coordinator.data.items() - for description in MOWER_SENSOR_TYPES - if description.exists_fn(data) - ) - - def _async_work_area_listener() -> None: - """Listen for new work areas and add sensor entities if they did not exist. - - Listening for deletable work areas is managed in the number platform. - """ - for mower_id in coordinator.data: - if ( - coordinator.data[mower_id].capabilities.work_areas - and (_work_areas := coordinator.data[mower_id].work_areas) is not None - ): - received_work_areas = set(_work_areas.keys()) - new_work_areas = received_work_areas - current_work_areas.get( - mower_id, set() + entities: list[SensorEntity] = [] + for mower_id in coordinator.data: + if coordinator.data[mower_id].capabilities.work_areas: + _work_areas = coordinator.data[mower_id].work_areas + if _work_areas is not None: + entities.extend( + WorkAreaSensorEntity( + mower_id, coordinator, description, work_area_id + ) + for description in WORK_AREA_SENSOR_TYPES + for work_area_id in _work_areas + if description.exists_fn(_work_areas[work_area_id]) ) - if new_work_areas: - current_work_areas.setdefault(mower_id, set()).update( - new_work_areas - ) - async_add_entities( - WorkAreaSensorEntity( - mower_id, coordinator, description, work_area_id - ) - for description in WORK_AREA_SENSOR_TYPES - for work_area_id in new_work_areas - if description.exists_fn(_work_areas[work_area_id]) - ) + entities.extend( + AutomowerSensorEntity(mower_id, coordinator, description) + for description in MOWER_SENSOR_TYPES + if description.exists_fn(coordinator.data[mower_id]) + ) + async_add_entities(entities) - coordinator.async_add_listener(_async_work_area_listener) - _async_work_area_listener() + def _async_add_new_work_areas(mower_id: str, work_area_ids: set[int]) -> None: + mower_data = coordinator.data[mower_id] + if mower_data.work_areas is None: + return + + async_add_entities( + WorkAreaSensorEntity(mower_id, coordinator, description, work_area_id) + for description in WORK_AREA_SENSOR_TYPES + for work_area_id in work_area_ids + if work_area_id in mower_data.work_areas + and description.exists_fn(mower_data.work_areas[work_area_id]) + ) + + def _async_add_new_devices(mower_ids: set[str]) -> None: + async_add_entities( + AutomowerSensorEntity(mower_id, coordinator, description) + for mower_id in mower_ids + for description in MOWER_SENSOR_TYPES + if description.exists_fn(coordinator.data[mower_id]) + ) + for mower_id in mower_ids: + mower_data = coordinator.data[mower_id] + if mower_data.capabilities.work_areas and mower_data.work_areas is not None: + _async_add_new_work_areas( + mower_id, + set(mower_data.work_areas.keys()), + ) + + coordinator.new_devices_callbacks.append(_async_add_new_devices) + coordinator.new_areas_callbacks.append(_async_add_new_work_areas) class AutomowerSensorEntity(AutomowerBaseEntity, SensorEntity): diff --git a/homeassistant/components/husqvarna_automower/switch.py b/homeassistant/components/husqvarna_automower/switch.py index 352b4c59ba1..b8004e17066 100644 --- a/homeassistant/components/husqvarna_automower/switch.py +++ b/homeassistant/components/husqvarna_automower/switch.py @@ -7,7 +7,6 @@ from aioautomower.model import MowerModes, StayOutZones, Zone from homeassistant.components.switch import SwitchEntity from homeassistant.core import HomeAssistant -from homeassistant.helpers import entity_registry as er from homeassistant.helpers.entity_platform import AddEntitiesCallback from . import AutomowerConfigEntry @@ -31,82 +30,63 @@ async def async_setup_entry( ) -> None: """Set up switch platform.""" coordinator = entry.runtime_data - current_work_areas: dict[str, set[int]] = {} - current_stay_out_zones: dict[str, set[str]] = {} - - async_add_entities( + entities: list[SwitchEntity] = [] + entities.extend( AutomowerScheduleSwitchEntity(mower_id, coordinator) for mower_id in coordinator.data ) - - def _async_work_area_listener() -> None: - """Listen for new work areas and add switch entities if they did not exist. - - Listening for deletable work areas is managed in the number platform. - """ - for mower_id in coordinator.data: - if ( - coordinator.data[mower_id].capabilities.work_areas - and (_work_areas := coordinator.data[mower_id].work_areas) is not None - ): - received_work_areas = set(_work_areas.keys()) - new_work_areas = received_work_areas - current_work_areas.get( - mower_id, set() + for mower_id in coordinator.data: + if coordinator.data[mower_id].capabilities.stay_out_zones: + _stay_out_zones = coordinator.data[mower_id].stay_out_zones + if _stay_out_zones is not None: + entities.extend( + StayOutZoneSwitchEntity(coordinator, mower_id, stay_out_zone_uid) + for stay_out_zone_uid in _stay_out_zones.zones ) - if new_work_areas: - current_work_areas.setdefault(mower_id, set()).update( - new_work_areas - ) - async_add_entities( - WorkAreaSwitchEntity(coordinator, mower_id, work_area_id) - for work_area_id in new_work_areas - ) + if coordinator.data[mower_id].capabilities.work_areas: + _work_areas = coordinator.data[mower_id].work_areas + if _work_areas is not None: + entities.extend( + WorkAreaSwitchEntity(coordinator, mower_id, work_area_id) + for work_area_id in _work_areas + ) + async_add_entities(entities) - def _remove_stay_out_zone_entities( - removed_stay_out_zones: set, mower_id: str + def _async_add_new_stay_out_zones( + mower_id: str, stay_out_zone_uids: set[str] ) -> None: - """Remove all unused stay-out zones for all platforms.""" - entity_reg = er.async_get(hass) - for entity_entry in er.async_entries_for_config_entry( - entity_reg, entry.entry_id - ): - for stay_out_zone_uid in removed_stay_out_zones: - if entity_entry.unique_id.startswith(f"{mower_id}_{stay_out_zone_uid}"): - entity_reg.async_remove(entity_entry.entity_id) + async_add_entities( + StayOutZoneSwitchEntity(coordinator, mower_id, zone_uid) + for zone_uid in stay_out_zone_uids + ) - def _async_stay_out_zone_listener() -> None: - """Listen for new stay-out zones and add/remove switch entities if they did not exist.""" - for mower_id in coordinator.data: + def _async_add_new_work_areas(mower_id: str, work_area_ids: set[int]) -> None: + async_add_entities( + WorkAreaSwitchEntity(coordinator, mower_id, work_area_id) + for work_area_id in work_area_ids + ) + + def _async_add_new_devices(mower_ids: set[str]) -> None: + async_add_entities( + AutomowerScheduleSwitchEntity(mower_id, coordinator) + for mower_id in mower_ids + ) + for mower_id in mower_ids: + mower_data = coordinator.data[mower_id] if ( - coordinator.data[mower_id].capabilities.stay_out_zones - and (_stay_out_zones := coordinator.data[mower_id].stay_out_zones) - is not None + mower_data.capabilities.stay_out_zones + and mower_data.stay_out_zones is not None + and mower_data.stay_out_zones.zones is not None ): - received_stay_out_zones = set(_stay_out_zones.zones) - current_stay_out_zones_set = current_stay_out_zones.get(mower_id, set()) - new_stay_out_zones = ( - received_stay_out_zones - current_stay_out_zones_set + _async_add_new_stay_out_zones( + mower_id, set(mower_data.stay_out_zones.zones.keys()) ) - removed_stay_out_zones = ( - current_stay_out_zones_set - received_stay_out_zones - ) - if new_stay_out_zones: - current_stay_out_zones.setdefault(mower_id, set()).update( - new_stay_out_zones - ) - async_add_entities( - StayOutZoneSwitchEntity( - coordinator, mower_id, stay_out_zone_uid - ) - for stay_out_zone_uid in new_stay_out_zones - ) - if removed_stay_out_zones: - _remove_stay_out_zone_entities(removed_stay_out_zones, mower_id) + if mower_data.capabilities.work_areas and mower_data.work_areas is not None: + _async_add_new_work_areas(mower_id, set(mower_data.work_areas.keys())) - coordinator.async_add_listener(_async_work_area_listener) - coordinator.async_add_listener(_async_stay_out_zone_listener) - _async_work_area_listener() - _async_stay_out_zone_listener() + coordinator.new_devices_callbacks.append(_async_add_new_devices) + coordinator.new_zones_callbacks.append(_async_add_new_stay_out_zones) + coordinator.new_areas_callbacks.append(_async_add_new_work_areas) class AutomowerScheduleSwitchEntity(AutomowerControlEntity, SwitchEntity): diff --git a/tests/components/husqvarna_automower/conftest.py b/tests/components/husqvarna_automower/conftest.py index 0202cec05b9..49994e4f3ae 100644 --- a/tests/components/husqvarna_automower/conftest.py +++ b/tests/components/husqvarna_automower/conftest.py @@ -58,6 +58,15 @@ def mock_values(mower_time_zone) -> dict[str, MowerAttributes]: ) +@pytest.fixture(name="values_one_mower") +def mock_values_one_mower(mower_time_zone) -> dict[str, MowerAttributes]: + """Fixture to set correct scope for the token.""" + return mower_list_to_dictionary_dataclass( + load_json_value_fixture("mower1.json", DOMAIN), + mower_time_zone, + ) + + @pytest.fixture def mock_config_entry(jwt: str, expires_at: int, scope: str) -> MockConfigEntry: """Return the default mocked config entry.""" @@ -119,3 +128,26 @@ def mock_automower_client(values) -> Generator[AsyncMock]: return_value=mock, ): yield mock + + +@pytest.fixture +def mock_automower_client_one_mower(values) -> Generator[AsyncMock]: + """Mock a Husqvarna Automower client.""" + + async def listen() -> None: + """Mock listen.""" + listen_block = asyncio.Event() + await listen_block.wait() + pytest.fail("Listen was not cancelled!") + + mock = AsyncMock(spec=AutomowerSession) + mock.auth = AsyncMock(side_effect=ClientWebSocketResponse) + mock.commands = AsyncMock(spec_set=_MowerCommands) + mock.get_status.return_value = values + mock.start_listening = AsyncMock(side_effect=listen) + + with patch( + "homeassistant.components.husqvarna_automower.AutomowerSession", + return_value=mock, + ): + yield mock diff --git a/tests/components/husqvarna_automower/test_init.py b/tests/components/husqvarna_automower/test_init.py index ae688571d2c..627cd065e79 100644 --- a/tests/components/husqvarna_automower/test_init.py +++ b/tests/components/husqvarna_automower/test_init.py @@ -227,32 +227,79 @@ async def test_coordinator_automatic_registry_cleanup( device_registry: dr.DeviceRegistry, entity_registry: er.EntityRegistry, values: dict[str, MowerAttributes], + freezer: FrozenDateTimeFactory, ) -> None: """Test automatic registry cleanup.""" await setup_integration(hass, mock_config_entry) entry = hass.config_entries.async_entries(DOMAIN)[0] await hass.async_block_till_done() + # Count current entitties and devices current_entites = len( er.async_entries_for_config_entry(entity_registry, entry.entry_id) ) current_devices = len( dr.async_entries_for_config_entry(device_registry, entry.entry_id) ) - - values.pop(TEST_MOWER_ID) + # Remove mower 2 and check if it worked + mower2 = values.pop("1234") mock_automower_client.get_status.return_value = values - await hass.config_entries.async_reload(mock_config_entry.entry_id) + freezer.tick(SCAN_INTERVAL) + async_fire_time_changed(hass) await hass.async_block_till_done() assert ( len(er.async_entries_for_config_entry(entity_registry, entry.entry_id)) - == current_entites - 37 + == current_entites - 12 ) assert ( len(dr.async_entries_for_config_entry(device_registry, entry.entry_id)) == current_devices - 1 ) + # Add mower 2 and check if it worked + values["1234"] = mower2 + mock_automower_client.get_status.return_value = values + freezer.tick(SCAN_INTERVAL) + async_fire_time_changed(hass) + await hass.async_block_till_done() + assert ( + len(er.async_entries_for_config_entry(entity_registry, entry.entry_id)) + == current_entites + ) + assert ( + len(dr.async_entries_for_config_entry(device_registry, entry.entry_id)) + == current_devices + ) + + # Remove mower 1 and check if it worked + mower1 = values.pop(TEST_MOWER_ID) + mock_automower_client.get_status.return_value = values + freezer.tick(SCAN_INTERVAL) + async_fire_time_changed(hass) + await hass.async_block_till_done() + + assert len(er.async_entries_for_config_entry(entity_registry, entry.entry_id)) == 12 + assert ( + len(dr.async_entries_for_config_entry(device_registry, entry.entry_id)) + == current_devices - 1 + ) + # Add mower 1 and check if it worked + values[TEST_MOWER_ID] = mower1 + mock_automower_client.get_status.return_value = values + freezer.tick(SCAN_INTERVAL) + async_fire_time_changed(hass) + await hass.async_block_till_done() + freezer.tick(SCAN_INTERVAL) + async_fire_time_changed(hass) + await hass.async_block_till_done() + assert ( + len(dr.async_entries_for_config_entry(device_registry, entry.entry_id)) + == current_devices + ) + assert ( + len(er.async_entries_for_config_entry(entity_registry, entry.entry_id)) + == current_entites + ) async def test_add_and_remove_work_area(