Automatic cleanup of entity and device registry in Tankerkoenig (#114573)

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

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.fuel_types: list[str] = self.config_entry.data[CONF_FUEL_TYPES]
-        self.show_on_map: bool = entry.options[CONF_SHOW_ON_MAP]
+        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. "

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