From 02b3da8f80aa4e2ad58e9fc8c814fc283d84ced8 Mon Sep 17 00:00:00 2001 From: Thomas55555 <59625598+Thomas55555@users.noreply.github.com> Date: Sun, 22 Sep 2024 16:06:01 +0200 Subject: [PATCH] Automatic device cleanup for Husqvarna Automower (#126384) * Automatic device cleanup for Husqvarna Automower * fix copy&paste mistake * typing * overwrite type in coordinator --- .../husqvarna_automower/__init__.py | 30 +++++++++++++++- .../husqvarna_automower/coordinator.py | 2 ++ .../husqvarna_automower/test_init.py | 36 +++++++++++++++++-- 3 files changed, 65 insertions(+), 3 deletions(-) diff --git a/homeassistant/components/husqvarna_automower/__init__.py b/homeassistant/components/husqvarna_automower/__init__.py index 6e987b679ed..117ded0dcf9 100644 --- a/homeassistant/components/husqvarna_automower/__init__.py +++ b/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 + ) diff --git a/homeassistant/components/husqvarna_automower/coordinator.py b/homeassistant/components/husqvarna_automower/coordinator.py index 817789727ca..458ff50dac9 100644 --- a/homeassistant/components/husqvarna_automower/coordinator.py +++ b/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: diff --git a/tests/components/husqvarna_automower/test_init.py b/tests/components/husqvarna_automower/test_init.py index 84fe1b9e891..ab80aea5a3f 100644 --- a/tests/components/husqvarna_automower/test_init.py +++ b/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