Automatic device cleanup for Husqvarna Automower (#126384)

* Automatic device cleanup for Husqvarna Automower

* fix copy&paste mistake

* typing

* overwrite type in coordinator

homeassistant/components/husqvarna_automower/__init__.py

@@ -9,9 +9,15 @@ 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
+from homeassistant.helpers import (
+    aiohttp_client,
+    config_entry_oauth2_flow,
+    device_registry as dr,
+    entity_registry as er,
+)
 
 from . import api
+from .const import DOMAIN
 from .coordinator import AutomowerDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
@@ -53,6 +59,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: AutomowerConfigEntry) ->
 
     coordinator = AutomowerDataUpdateCoordinator(hass, automower_api, entry)
     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(
@@ -73,3 +81,23 @@ 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: ConfigEntry, available_devices: list[str]
+) -> None:
+    """Cleanup entity and device registry from removed devices."""
+    entity_reg = er.async_get(hass)
+    for entity in er.async_entries_for_config_entry(entity_reg, config_entry.entry_id):
+        if entity.unique_id.split("_")[0] not in available_devices:
+            _LOGGER.debug("Removing obsolete entity entry %s", entity.entity_id)
+            entity_reg.async_remove(entity.entity_id)
+
+    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
+            )

homeassistant/components/husqvarna_automower/coordinator.py

@@ -27,6 +27,8 @@ SCAN_INTERVAL = timedelta(minutes=8)
 class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[dict[str, MowerAttributes]]):
     """Class to manage fetching Husqvarna data."""
 
+    config_entry: ConfigEntry
+
     def __init__(
         self, hass: HomeAssistant, api: AutomowerSession, entry: ConfigEntry
     ) -> None:

tests/components/husqvarna_automower/test_init.py

@@ -10,6 +10,7 @@ from aioautomower.exceptions import (
     AuthException,
     HusqvarnaWSServerHandshakeError,
 )
+from aioautomower.utils import mower_list_to_dictionary_dataclass
 from freezegun.api import FrozenDateTimeFactory
 import pytest
 from syrupy.assertion import SnapshotAssertion
@@ -17,12 +18,16 @@ from syrupy.assertion import SnapshotAssertion
 from homeassistant.components.husqvarna_automower.const import DOMAIN, OAUTH2_TOKEN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
 from . import setup_integration
 from .const import TEST_MOWER_ID
 
-from tests.common import MockConfigEntry, async_fire_time_changed
+from tests.common import (
+    MockConfigEntry,
+    async_fire_time_changed,
+    load_json_value_fixture,
+)
 from tests.test_util.aiohttp import AiohttpClientMocker
 
 
@@ -160,3 +165,30 @@ async def test_device_info(
         identifiers={(DOMAIN, TEST_MOWER_ID)},
     )
     assert reg_device == snapshot
+
+
+async def test_coordinator_automatic_registry_cleanup(
+    hass: HomeAssistant,
+    mock_automower_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+) -> 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()
+
+    assert len(er.async_entries_for_config_entry(entity_registry, entry.entry_id)) == 42
+    assert len(dr.async_entries_for_config_entry(device_registry, entry.entry_id)) == 2
+
+    values = mower_list_to_dictionary_dataclass(
+        load_json_value_fixture("mower.json", DOMAIN)
+    )
+    values.pop(TEST_MOWER_ID)
+    mock_automower_client.get_status.return_value = values
+    await hass.config_entries.async_reload(mock_config_entry.entry_id)
+    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)) == 1