Automatic cleanup of entity and device registry in AVM FRITZ!SmartHome (#114601)

This commit is contained in:
Michael 2024-04-20 12:13:56 +02:00 committed by GitHub
parent 0ea1564248
commit 354c20a57b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 165 additions and 75 deletions

View File

@ -51,12 +51,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
has_templates = await hass.async_add_executor_job(fritz.has_templates)
LOGGER.debug("enable smarthome templates: %s", has_templates)
coordinator = FritzboxDataUpdateCoordinator(hass, entry, has_templates)
await coordinator.async_config_entry_first_refresh()
hass.data[DOMAIN][entry.entry_id][CONF_COORDINATOR] = coordinator
def _update_unique_id(entry: RegistryEntry) -> dict[str, str] | None:
"""Update unique ID of entity entry."""
if (
@ -79,6 +73,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
await async_migrate_entries(hass, entry.entry_id, _update_unique_id)
coordinator = FritzboxDataUpdateCoordinator(hass, entry.entry_id, has_templates)
await coordinator.async_setup()
hass.data[DOMAIN][entry.entry_id][CONF_COORDINATOR] = coordinator
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
def logout_fritzbox(event: Event) -> None:

View File

@ -12,6 +12,7 @@ from requests.exceptions import ConnectionError as RequestConnectionError, HTTPE
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryAuthFailed
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
from .const import CONF_CONNECTIONS, DOMAIN, LOGGER
@ -28,27 +29,55 @@ class FritzboxCoordinatorData:
class FritzboxDataUpdateCoordinator(DataUpdateCoordinator[FritzboxCoordinatorData]):
"""Fritzbox Smarthome device data update coordinator."""
config_entry: ConfigEntry
configuration_url: str
def __init__(
self, hass: HomeAssistant, entry: ConfigEntry, has_templates: bool
) -> None:
def __init__(self, hass: HomeAssistant, name: str, has_templates: bool) -> None:
"""Initialize the Fritzbox Smarthome device coordinator."""
self.entry = entry
self.fritz: Fritzhome = hass.data[DOMAIN][self.entry.entry_id][CONF_CONNECTIONS]
super().__init__(
hass,
LOGGER,
name=name,
update_interval=timedelta(seconds=30),
)
self.fritz: Fritzhome = hass.data[DOMAIN][self.config_entry.entry_id][
CONF_CONNECTIONS
]
self.configuration_url = self.fritz.get_prefixed_host()
self.has_templates = has_templates
self.new_devices: set[str] = set()
self.new_templates: set[str] = set()
super().__init__(
hass,
LOGGER,
name=entry.entry_id,
update_interval=timedelta(seconds=30),
self.data = FritzboxCoordinatorData({}, {})
async def async_setup(self) -> None:
"""Set up the coordinator."""
await self.async_config_entry_first_refresh()
self.cleanup_removed_devices(
list(self.data.devices) + list(self.data.templates)
)
self.data = FritzboxCoordinatorData({}, {})
def cleanup_removed_devices(self, avaiable_ains: list[str]) -> None:
"""Cleanup entity and device registry from removed devices."""
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 avaiable_ains:
LOGGER.debug("Removing obsolete entity entry %s", entity.entity_id)
entity_reg.async_remove(entity.entity_id)
device_reg = dr.async_get(self.hass)
identifiers = {(DOMAIN, ain) for ain in avaiable_ains}
for device in dr.async_entries_for_config_entry(
device_reg, self.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=self.config_entry.entry_id
)
def _update_fritz_devices(self) -> FritzboxCoordinatorData:
"""Update all fritzbox device data."""
@ -95,6 +124,12 @@ class FritzboxDataUpdateCoordinator(DataUpdateCoordinator[FritzboxCoordinatorDat
self.new_devices = device_data.keys() - self.data.devices.keys()
self.new_templates = template_data.keys() - self.data.templates.keys()
if (
self.data.devices.keys() - device_data.keys()
or self.data.templates.keys() - template_data.keys()
):
self.cleanup_removed_devices(list(device_data) + list(template_data))
return FritzboxCoordinatorData(devices=device_data, templates=template_data)
async def _async_update_data(self) -> FritzboxCoordinatorData:

View File

@ -0,0 +1,111 @@
"""Tests for the AVM Fritz!Box integration."""
from __future__ import annotations
from datetime import timedelta
from unittest.mock import Mock
from pyfritzhome import LoginError
from requests.exceptions import ConnectionError, HTTPError
from homeassistant.components.fritzbox.const import DOMAIN as FB_DOMAIN
from homeassistant.config_entries import ConfigEntryState
from homeassistant.const import CONF_DEVICES
from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.util.dt import utcnow
from . import FritzDeviceCoverMock, FritzDeviceSwitchMock
from .const import MOCK_CONFIG
from tests.common import MockConfigEntry, async_fire_time_changed
async def test_coordinator_update_after_reboot(
hass: HomeAssistant, fritz: Mock
) -> None:
"""Test coordinator after reboot."""
entry = MockConfigEntry(
domain=FB_DOMAIN,
data=MOCK_CONFIG[FB_DOMAIN][CONF_DEVICES][0],
unique_id="any",
)
entry.add_to_hass(hass)
fritz().update_devices.side_effect = [HTTPError(), ""]
assert await hass.config_entries.async_setup(entry.entry_id)
assert fritz().update_devices.call_count == 2
assert fritz().update_templates.call_count == 1
assert fritz().get_devices.call_count == 1
assert fritz().get_templates.call_count == 1
assert fritz().login.call_count == 2
async def test_coordinator_update_after_password_change(
hass: HomeAssistant, fritz: Mock
) -> None:
"""Test coordinator after password change."""
entry = MockConfigEntry(
domain=FB_DOMAIN,
data=MOCK_CONFIG[FB_DOMAIN][CONF_DEVICES][0],
unique_id="any",
)
entry.add_to_hass(hass)
fritz().update_devices.side_effect = HTTPError()
fritz().login.side_effect = ["", LoginError("some_user")]
assert not await hass.config_entries.async_setup(entry.entry_id)
assert fritz().update_devices.call_count == 1
assert fritz().get_devices.call_count == 0
assert fritz().get_templates.call_count == 0
assert fritz().login.call_count == 2
async def test_coordinator_update_when_unreachable(
hass: HomeAssistant, fritz: Mock
) -> None:
"""Test coordinator after reboot."""
entry = MockConfigEntry(
domain=FB_DOMAIN,
data=MOCK_CONFIG[FB_DOMAIN][CONF_DEVICES][0],
unique_id="any",
)
entry.add_to_hass(hass)
fritz().update_devices.side_effect = [ConnectionError(), ""]
assert not await hass.config_entries.async_setup(entry.entry_id)
assert entry.state is ConfigEntryState.SETUP_RETRY
async def test_coordinator_automatic_registry_cleanup(
hass: HomeAssistant,
fritz: Mock,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Test automatic registry cleanup."""
fritz().get_devices.return_value = [
FritzDeviceSwitchMock(ain="fake ain switch", name="fake_switch"),
FritzDeviceCoverMock(ain="fake ain cover", name="fake_cover"),
]
entry = MockConfigEntry(
domain=FB_DOMAIN,
data=MOCK_CONFIG[FB_DOMAIN][CONF_DEVICES][0],
unique_id="any",
)
entry.add_to_hass(hass)
await hass.config_entries.async_setup(entry.entry_id)
await hass.async_block_till_done(wait_background_tasks=True)
assert len(er.async_entries_for_config_entry(entity_registry, entry.entry_id)) == 11
assert len(dr.async_entries_for_config_entry(device_registry, entry.entry_id)) == 2
fritz().get_devices.return_value = [
FritzDeviceSwitchMock(ain="fake ain switch", name="fake_switch")
]
async_fire_time_changed(hass, utcnow() + timedelta(seconds=35))
await hass.async_block_till_done(wait_background_tasks=True)
assert len(er.async_entries_for_config_entry(entity_registry, entry.entry_id)) == 8
assert len(dr.async_entries_for_config_entry(device_registry, entry.entry_id)) == 1

View File

@ -6,7 +6,7 @@ from unittest.mock import Mock, call, patch
from pyfritzhome import LoginError
import pytest
from requests.exceptions import ConnectionError, HTTPError
from requests.exceptions import ConnectionError as RequestConnectionError
from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR_DOMAIN
from homeassistant.components.fritzbox.const import DOMAIN as FB_DOMAIN
@ -80,6 +80,7 @@ async def test_update_unique_id(
new_unique_id: str,
) -> None:
"""Test unique_id update of integration."""
fritz().get_devices.return_value = [FritzDeviceSwitchMock()]
entry = MockConfigEntry(
domain=FB_DOMAIN,
data=MOCK_CONFIG[FB_DOMAIN][CONF_DEVICES][0],
@ -138,6 +139,7 @@ async def test_update_unique_id_no_change(
unique_id: str,
) -> None:
"""Test unique_id is not updated of integration."""
fritz().get_devices.return_value = [FritzDeviceSwitchMock()]
entry = MockConfigEntry(
domain=FB_DOMAIN,
data=MOCK_CONFIG[FB_DOMAIN][CONF_DEVICES][0],
@ -158,62 +160,6 @@ async def test_update_unique_id_no_change(
assert entity_migrated.unique_id == unique_id
async def test_coordinator_update_after_reboot(
hass: HomeAssistant, fritz: Mock
) -> None:
"""Test coordinator after reboot."""
entry = MockConfigEntry(
domain=FB_DOMAIN,
data=MOCK_CONFIG[FB_DOMAIN][CONF_DEVICES][0],
unique_id="any",
)
entry.add_to_hass(hass)
fritz().update_devices.side_effect = [HTTPError(), ""]
assert await hass.config_entries.async_setup(entry.entry_id)
assert fritz().update_devices.call_count == 2
assert fritz().update_templates.call_count == 1
assert fritz().get_devices.call_count == 1
assert fritz().get_templates.call_count == 1
assert fritz().login.call_count == 2
async def test_coordinator_update_after_password_change(
hass: HomeAssistant, fritz: Mock
) -> None:
"""Test coordinator after password change."""
entry = MockConfigEntry(
domain=FB_DOMAIN,
data=MOCK_CONFIG[FB_DOMAIN][CONF_DEVICES][0],
unique_id="any",
)
entry.add_to_hass(hass)
fritz().update_devices.side_effect = HTTPError()
fritz().login.side_effect = ["", LoginError("some_user")]
assert not await hass.config_entries.async_setup(entry.entry_id)
assert fritz().update_devices.call_count == 1
assert fritz().get_devices.call_count == 0
assert fritz().get_templates.call_count == 0
assert fritz().login.call_count == 2
async def test_coordinator_update_when_unreachable(
hass: HomeAssistant, fritz: Mock
) -> None:
"""Test coordinator after reboot."""
entry = MockConfigEntry(
domain=FB_DOMAIN,
data=MOCK_CONFIG[FB_DOMAIN][CONF_DEVICES][0],
unique_id="any",
)
entry.add_to_hass(hass)
fritz().update_devices.side_effect = [ConnectionError(), ""]
assert not await hass.config_entries.async_setup(entry.entry_id)
assert entry.state is ConfigEntryState.SETUP_RETRY
async def test_unload_remove(hass: HomeAssistant, fritz: Mock) -> None:
"""Test unload and remove of integration."""
fritz().get_devices.return_value = [FritzDeviceSwitchMock()]
@ -325,7 +271,7 @@ async def test_raise_config_entry_not_ready_when_offline(hass: HomeAssistant) ->
entry.add_to_hass(hass)
with patch(
"homeassistant.components.fritzbox.Fritzhome.login",
side_effect=ConnectionError(),
side_effect=RequestConnectionError(),
) as mock_login:
await hass.config_entries.async_setup(entry.entry_id)
await hass.async_block_till_done()