From 2e11a61726a29ee0720238d8b081755aa4d94c4e Mon Sep 17 00:00:00 2001 From: Michael <35783820+mib1185@users.noreply.github.com> Date: Mon, 1 Apr 2024 16:04:18 +0200 Subject: [PATCH] Automatic cleanup of entity and device registry in Tankerkoenig (#114573) --- .../components/tankerkoenig/__init__.py | 1 - .../components/tankerkoenig/coordinator.py | 34 +++++++-- .../tankerkoenig/test_coordinator.py | 71 ++++++++++++++++++- 3 files changed, 98 insertions(+), 8 deletions(-) diff --git a/homeassistant/components/tankerkoenig/__init__.py b/homeassistant/components/tankerkoenig/__init__.py index 7443fa72b5b..ac009b7a274 100644 --- a/homeassistant/components/tankerkoenig/__init__.py +++ b/homeassistant/components/tankerkoenig/__init__.py @@ -21,7 +21,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: coordinator = TankerkoenigDataUpdateCoordinator( hass, - entry, name=entry.unique_id or DOMAIN, update_interval=DEFAULT_SCAN_INTERVAL, ) diff --git a/homeassistant/components/tankerkoenig/coordinator.py b/homeassistant/components/tankerkoenig/coordinator.py index 447099d2dca..b7a45b65d0a 100644 --- a/homeassistant/components/tankerkoenig/coordinator.py +++ b/homeassistant/components/tankerkoenig/coordinator.py @@ -17,9 +17,10 @@ from aiotankerkoenig import ( ) from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_API_KEY, CONF_SHOW_ON_MAP +from homeassistant.const import ATTR_ID, CONF_API_KEY, CONF_SHOW_ON_MAP from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady +from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed @@ -31,10 +32,11 @@ _LOGGER = logging.getLogger(__name__) class TankerkoenigDataUpdateCoordinator(DataUpdateCoordinator): """Get the latest data from the API.""" + config_entry: ConfigEntry + def __init__( self, hass: HomeAssistant, - entry: ConfigEntry, name: str, update_interval: int, ) -> None: @@ -47,13 +49,14 @@ class TankerkoenigDataUpdateCoordinator(DataUpdateCoordinator): update_interval=timedelta(minutes=update_interval), ) - self._selected_stations: list[str] = entry.data[CONF_STATIONS] + self._selected_stations: list[str] = self.config_entry.data[CONF_STATIONS] self.stations: dict[str, Station] = {} - self.fuel_types: list[str] = entry.data[CONF_FUEL_TYPES] - self.show_on_map: bool = entry.options[CONF_SHOW_ON_MAP] + self.fuel_types: list[str] = self.config_entry.data[CONF_FUEL_TYPES] + self.show_on_map: bool = self.config_entry.options[CONF_SHOW_ON_MAP] self._tankerkoenig = Tankerkoenig( - api_key=entry.data[CONF_API_KEY], session=async_get_clientsession(hass) + api_key=self.config_entry.data[CONF_API_KEY], + session=async_get_clientsession(hass), ) async def async_setup(self) -> None: @@ -81,6 +84,25 @@ class TankerkoenigDataUpdateCoordinator(DataUpdateCoordinator): self.stations[station_id] = station + entity_reg = er.async_get(self.hass) + for entity in er.async_entries_for_config_entry( + entity_reg, self.config_entry.entry_id + ): + if entity.unique_id.split("_")[0] not in self._selected_stations: + _LOGGER.debug("Removing obsolete entity entry %s", entity.entity_id) + entity_reg.async_remove(entity.entity_id) + + device_reg = dr.async_get(self.hass) + for device in dr.async_entries_for_config_entry( + device_reg, self.config_entry.entry_id + ): + if not any( + (ATTR_ID, station_id) in device.identifiers + for station_id in self._selected_stations + ): + _LOGGER.debug("Removing obsolete device entry %s", device.name) + device_reg.async_remove_device(device.id) + if len(self.stations) > 10: _LOGGER.warning( "Found more than 10 stations to check. " diff --git a/tests/components/tankerkoenig/test_coordinator.py b/tests/components/tankerkoenig/test_coordinator.py index 5a33cb95dd9..de65cd921be 100644 --- a/tests/components/tankerkoenig/test_coordinator.py +++ b/tests/components/tankerkoenig/test_coordinator.py @@ -13,10 +13,13 @@ from aiotankerkoenig.exceptions import ( ) import pytest +from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR_DOMAIN +from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN from homeassistant.components.tankerkoenig.const import DEFAULT_SCAN_INTERVAL, DOMAIN from homeassistant.config_entries import ConfigEntryState -from homeassistant.const import STATE_UNAVAILABLE +from homeassistant.const import ATTR_ID, STATE_UNAVAILABLE from homeassistant.core import HomeAssistant +from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.setup import async_setup_component import homeassistant.util.dt as dt_util @@ -121,3 +124,69 @@ async def test_setup_exception_logging( await hass.async_block_till_done() assert expected_log in caplog.text + + +async def test_automatic_registry_cleanup( + hass: HomeAssistant, + config_entry: MockConfigEntry, + tankerkoenig: AsyncMock, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, +) -> None: + """Test automatic registry cleanup for obsolete entity and devices entries.""" + # setup normal + config_entry.add_to_hass(hass) + assert await async_setup_component(hass, DOMAIN, {}) + await hass.async_block_till_done() + + assert ( + len(er.async_entries_for_config_entry(entity_registry, config_entry.entry_id)) + == 4 + ) + assert ( + len(dr.async_entries_for_config_entry(device_registry, config_entry.entry_id)) + == 1 + ) + + # add obsolete entity and device entries + obsolete_station_id = "aabbccddee-xxxx-xxxx-xxxx-ff11223344" + + entity_registry.async_get_or_create( + DOMAIN, + BINARY_SENSOR_DOMAIN, + f"{obsolete_station_id}_status", + config_entry=config_entry, + ) + entity_registry.async_get_or_create( + DOMAIN, + SENSOR_DOMAIN, + f"{obsolete_station_id}_e10", + config_entry=config_entry, + ) + device_registry.async_get_or_create( + config_entry_id=config_entry.entry_id, + identifiers={(ATTR_ID, obsolete_station_id)}, + name="Obsolete Station", + ) + + assert ( + len(er.async_entries_for_config_entry(entity_registry, config_entry.entry_id)) + == 6 + ) + assert ( + len(dr.async_entries_for_config_entry(device_registry, config_entry.entry_id)) + == 2 + ) + + # reload config entry to trigger automatic cleanup + await hass.config_entries.async_reload(config_entry.entry_id) + await hass.async_block_till_done() + + assert ( + len(er.async_entries_for_config_entry(entity_registry, config_entry.entry_id)) + == 4 + ) + assert ( + len(dr.async_entries_for_config_entry(device_registry, config_entry.entry_id)) + == 1 + )