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

This commit is contained in:
Michael 2024-04-01 16:04:18 +02:00 committed by GitHub
parent 7f9ad140f9
commit 2e11a61726
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 98 additions and 8 deletions

View File

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

View File

@ -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. "

View File

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