Dynamic devices for Husqvarna Automower (#133227)

* Dynamic devices for Husqvarna Automower

* callbacks

* add stayout-zones together

* add alltogether on init

* fix stale lock names

* also for workareas

* separate "normal" vs callback entity adding

* mark quality scale

* Apply suggestions from code review

Co-authored-by: Josef Zweck <josef@zweck.dev>

* Apply suggestions from code review

Co-authored-by: Josef Zweck <josef@zweck.dev>

---------

Co-authored-by: Josef Zweck <josef@zweck.dev>

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 (
+from homeassistant.helpers import aiohttp_client, config_entry_oauth2_flow
-    aiohttp_client,
-    config_entry_oauth2_flow,
-    device_registry as dr,
-    entity_registry as er,
-)
 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)

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)
+    def _async_add_new_devices(mower_ids: set[str]) -> None:
-        for mower_id in coordinator.data
+        async_add_entities(
-        for description in MOWER_BINARY_SENSOR_TYPES
+            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):

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)
+    def _async_add_new_devices(mower_ids: set[str]) -> None:
-        for mower_id in coordinator.data
+        async_add_entities(
-        for description in MOWER_BUTTON_TYPES
+            AutomowerButtonEntity(mower_id, coordinator, description)
-        if description.exists_fn(coordinator.data[mower_id])
+            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):

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):

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

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)
+    def _async_add_new_devices(mower_ids: set[str]) -> None:
-        for mower_id in coordinator.data
+        async_add_entities(
-        if coordinator.data[mower_id].capabilities.position
+            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):

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",

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]] = {}
+    entities: list[NumberEntity] = []
-
+    for mower_id in coordinator.data:
-    async_add_entities(
+        if coordinator.data[mower_id].capabilities.work_areas:
-        AutomowerNumberEntity(mower_id, coordinator, description)
+            _work_areas = coordinator.data[mower_id].work_areas
-        for mower_id in coordinator.data
+            if _work_areas is not None:
-        for description in MOWER_NUMBER_TYPES
+                entities.extend(
-        if description.exists_fn(coordinator.data[mower_id])
+                    WorkAreaNumberEntity(
-    )
+                        mower_id, coordinator, description, work_area_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
                     )
+                    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:
+    def _async_add_new_work_areas(mower_id: str, work_area_ids: set[int]) -> None:
-                    remove_work_area_entities(hass, entry, removed_work_areas, mower_id)
+        async_add_entities(
-                    current_work_area_set.difference_update(removed_work_areas)
+            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)
+    def _async_add_new_devices(mower_ids: set[str]) -> None:
-    _async_work_area_listener()
+        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):

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:
+  dynamic-devices: done
-    status: todo
-    comment: Add devices dynamically
   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:
+  stale-devices: done
-    status: todo
-    comment: We only remove devices on reload
 
   # Platinum
   async-dependency: done

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)
+    def _async_add_new_devices(mower_ids: set[str]) -> None:
-        for mower_id in coordinator.data
+        async_add_entities(
-        if coordinator.data[mower_id].capabilities.headlights
+            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):

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]] = {}
+    entities: list[SensorEntity] = []
-
+    for mower_id in coordinator.data:
-    async_add_entities(
+        if coordinator.data[mower_id].capabilities.work_areas:
-        AutomowerSensorEntity(mower_id, coordinator, description)
+            _work_areas = coordinator.data[mower_id].work_areas
-        for mower_id, data in coordinator.data.items()
+            if _work_areas is not None:
-        for description in MOWER_SENSOR_TYPES
+                entities.extend(
-        if description.exists_fn(data)
+                    WorkAreaSensorEntity(
-    )
+                        mower_id, coordinator, description, work_area_id
-
+                    )
-    def _async_work_area_listener() -> None:
+                    for description in WORK_AREA_SENSOR_TYPES
-        """Listen for new work areas and add sensor entities if they did not exist.
+                    for work_area_id in _work_areas
-
+                    if description.exists_fn(_work_areas[work_area_id])
-        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()
                 )
-                if new_work_areas:
+        entities.extend(
-                    current_work_areas.setdefault(mower_id, set()).update(
+            AutomowerSensorEntity(mower_id, coordinator, description)
-                        new_work_areas
+            for description in MOWER_SENSOR_TYPES
-                    )
+            if description.exists_fn(coordinator.data[mower_id])
-                    async_add_entities(
+        )
-                        WorkAreaSensorEntity(
+    async_add_entities(entities)
-                            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])
-                    )
 
-    coordinator.async_add_listener(_async_work_area_listener)
+    def _async_add_new_work_areas(mower_id: str, work_area_ids: set[int]) -> None:
-    _async_work_area_listener()
+        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):

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]] = {}
+    entities: list[SwitchEntity] = []
-    current_stay_out_zones: dict[str, set[str]] = {}
+    entities.extend(
-
-    async_add_entities(
         AutomowerScheduleSwitchEntity(mower_id, coordinator)
         for mower_id in coordinator.data
     )
-
+    for mower_id in coordinator.data:
-    def _async_work_area_listener() -> None:
+        if coordinator.data[mower_id].capabilities.stay_out_zones:
-        """Listen for new work areas and add switch entities if they did not exist.
+            _stay_out_zones = coordinator.data[mower_id].stay_out_zones
-
+            if _stay_out_zones is not None:
-        Listening for deletable work areas is managed in the number platform.
+                entities.extend(
-        """
+                    StayOutZoneSwitchEntity(coordinator, mower_id, stay_out_zone_uid)
-        for mower_id in coordinator.data:
+                    for stay_out_zone_uid in _stay_out_zones.zones
-            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()
                 )
-                if new_work_areas:
+        if coordinator.data[mower_id].capabilities.work_areas:
-                    current_work_areas.setdefault(mower_id, set()).update(
+            _work_areas = coordinator.data[mower_id].work_areas
-                        new_work_areas
+            if _work_areas is not None:
-                    )
+                entities.extend(
-                    async_add_entities(
+                    WorkAreaSwitchEntity(coordinator, mower_id, work_area_id)
-                        WorkAreaSwitchEntity(coordinator, mower_id, work_area_id)
+                    for work_area_id in _work_areas
-                        for work_area_id in new_work_areas
+                )
-                    )
+    async_add_entities(entities)
 
-    def _remove_stay_out_zone_entities(
+    def _async_add_new_stay_out_zones(
-        removed_stay_out_zones: set, mower_id: str
+        mower_id: str, stay_out_zone_uids: set[str]
     ) -> None:
-        """Remove all unused stay-out zones for all platforms."""
+        async_add_entities(
-        entity_reg = er.async_get(hass)
+            StayOutZoneSwitchEntity(coordinator, mower_id, zone_uid)
-        for entity_entry in er.async_entries_for_config_entry(
+            for zone_uid in stay_out_zone_uids
-            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)
 
-    def _async_stay_out_zone_listener() -> None:
+    def _async_add_new_work_areas(mower_id: str, work_area_ids: set[int]) -> None:
-        """Listen for new stay-out zones and add/remove switch entities if they did not exist."""
+        async_add_entities(
-        for mower_id in coordinator.data:
+            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
+                mower_data.capabilities.stay_out_zones
-                and (_stay_out_zones := coordinator.data[mower_id].stay_out_zones)
+                and mower_data.stay_out_zones is not None
-                is not None
+                and mower_data.stay_out_zones.zones is not None
             ):
-                received_stay_out_zones = set(_stay_out_zones.zones)
+                _async_add_new_stay_out_zones(
-                current_stay_out_zones_set = current_stay_out_zones.get(mower_id, set())
+                    mower_id, set(mower_data.stay_out_zones.zones.keys())
-                new_stay_out_zones = (
-                    received_stay_out_zones - current_stay_out_zones_set
                 )
-                removed_stay_out_zones = (
+            if mower_data.capabilities.work_areas and mower_data.work_areas is not None:
-                    current_stay_out_zones_set - received_stay_out_zones
+                _async_add_new_work_areas(mower_id, set(mower_data.work_areas.keys()))
-                )
-                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)
 
-    coordinator.async_add_listener(_async_work_area_listener)
+    coordinator.new_devices_callbacks.append(_async_add_new_devices)
-    coordinator.async_add_listener(_async_stay_out_zone_listener)
+    coordinator.new_zones_callbacks.append(_async_add_new_stay_out_zones)
-    _async_work_area_listener()
+    coordinator.new_areas_callbacks.append(_async_add_new_work_areas)
-    _async_stay_out_zone_listener()
 
 
 class AutomowerScheduleSwitchEntity(AutomowerControlEntity, SwitchEntity):

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

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)
     )
-
+    # Remove mower 2 and check if it worked
-    values.pop(TEST_MOWER_ID)
+    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(