Improve removal of stale entities/devices in Husqvarna Automower (home-assistant#148428)

Co-authored-by: Abílio Costa <[email protected]>

Author: Thomas55555

homeassistant/components/husqvarna_automower/coordinator.py

@@ -58,9 +58,6 @@ def __init__(
         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]] = {}
 
     @override
     @callback
@@ -87,11 +84,15 @@ def _on_data_update(self) -> None:
         """Handle data updates and process dynamic entity management."""
         if self.data is not None:
             self._async_add_remove_devices()
-            for mower_id in self.data:
-                if self.data[mower_id].capabilities.stay_out_zones:
-                    self._async_add_remove_stay_out_zones()
-                if self.data[mower_id].capabilities.work_areas:
-                    self._async_add_remove_work_areas()
+            if any(
+                mower_data.capabilities.stay_out_zones
+                for mower_data in self.data.values()
+            ):
+                self._async_add_remove_stay_out_zones()
+            if any(
+                mower_data.capabilities.work_areas for mower_data in self.data.values()
+            ):
+                self._async_add_remove_work_areas()
 
     @callback
     def handle_websocket_updates(self, ws_data: MowerDictionary) -> None:
@@ -161,44 +162,36 @@ async def client_listen(
             )
 
     def _async_add_remove_devices(self) -> None:
-        """Add new device, remove non-existing device."""
+        """Add new devices and remove orphaned devices from the registry."""
         current_devices = set(self.data)
+        device_registry = dr.async_get(self.hass)
 
-        # Skip update if no changes
-        if current_devices == self._devices_last_update:
-            return
+        registered_devices: set[str] = {
+            str(mower_id)
+            for device in device_registry.devices.get_devices_for_config_entry_id(
+                self.config_entry.entry_id
+            )
+            for domain, mower_id in device.identifiers
+            if domain == DOMAIN
+        }
 
-        # 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)
+        orphaned_devices = registered_devices - current_devices
+        if orphaned_devices:
+            _LOGGER.debug("Removing orphaned devices: %s", orphaned_devices)
+            device_registry = dr.async_get(self.hass)
+            for mower_id in orphaned_devices:
+                dev = device_registry.async_get_device(identifiers={(DOMAIN, mower_id)})
+                if dev is not None:
+                    device_registry.async_update_device(
+                        device_id=dev.id,
+                        remove_config_entry_id=self.config_entry.entry_id,
+                    )
 
-        # Process new device
-        new_devices = current_devices - self._devices_last_update
+        new_devices = current_devices - registered_devices
         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)
+            _LOGGER.debug("New devices found: %s", new_devices)
+            for mower_callback in self.new_devices_callbacks:
+                mower_callback(new_devices)
 
     def _async_add_remove_stay_out_zones(self) -> None:
         """Add new stay-out zones, remove non-existing stay-out zones."""
@@ -209,42 +202,39 @@ def _async_add_remove_stay_out_zones(self) -> None:
             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
+        entity_registry = er.async_get(self.hass)
+        entries = er.async_entries_for_config_entry(
+            entity_registry, self.config_entry.entry_id
+        )
 
-        self._zones_last_update = self._update_stay_out_zones(current_zones)
+        registered_zones: dict[str, set[str]] = {}
+        for mower_id in self.data:
+            registered_zones[mower_id] = set()
+            for entry in entries:
+                uid = entry.unique_id
+                if uid.startswith(f"{mower_id}_") and uid.endswith("_stay_out_zones"):
+                    zone_id = uid.removeprefix(f"{mower_id}_").removesuffix(
+                        "_stay_out_zones"
+                    )
+                    registered_zones[mower_id].add(zone_id)
 
-    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, current_ids in current_zones.items():
+            known_ids = registered_zones.get(mower_id, set())
 
-        for mower_id, zones in new_zones.items():
-            for zone_callback in self.new_zones_callbacks:
-                zone_callback(mower_id, set(zones))
+            new_zones = current_ids - known_ids
+            removed_zones = known_ids - current_ids
 
-        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)
+            if new_zones:
+                _LOGGER.debug("New stay-out zones: %s", new_zones)
+                for zone_callback in self.new_zones_callbacks:
+                    zone_callback(mower_id, new_zones)
 
-        return current_zones
+            if removed_zones:
+                _LOGGER.debug("Removing stay-out zones: %s", removed_zones)
+                for entry in entries:
+                    for zone_id in removed_zones:
+                        if entry.unique_id == f"{mower_id}_{zone_id}_stay_out_zones":
+                            entity_registry.async_remove(entry.entity_id)
 
     def _async_add_remove_work_areas(self) -> None:
         """Add new work areas, remove non-existing work areas."""
@@ -254,39 +244,36 @@ def _async_add_remove_work_areas(self) -> None:
             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)
+        entries = er.async_entries_for_config_entry(
+            entity_registry, self.config_entry.entry_id
+        )
 
-        return current_areas
+        registered_areas: dict[str, set[int]] = {}
+        for mower_id in self.data:
+            registered_areas[mower_id] = set()
+            for entry in entries:
+                uid = entry.unique_id
+                if uid.startswith(f"{mower_id}_") and uid.endswith("_work_area"):
+                    parts = uid.removeprefix(f"{mower_id}_").split("_")
+                    area_id_str = parts[0] if parts else None
+                    if area_id_str and area_id_str.isdigit():
+                        registered_areas[mower_id].add(int(area_id_str))
+
+        for mower_id, current_ids in current_areas.items():
+            known_ids = registered_areas.get(mower_id, set())
+
+            new_areas = current_ids - known_ids
+            removed_areas = known_ids - current_ids
+
+            if new_areas:
+                _LOGGER.debug("New work areas: %s", new_areas)
+                for area_callback in self.new_areas_callbacks:
+                    area_callback(mower_id, new_areas)
+
+            if removed_areas:
+                _LOGGER.debug("Removing work areas: %s", removed_areas)
+                for entry in entries:
+                    for area_id in removed_areas:
+                        if entry.unique_id.startswith(f"{mower_id}_{area_id}_"):
+                            entity_registry.async_remove(entry.entity_id)