From 354c20a57b4e513a42a4b0b45fec403dc63d79b6 Mon Sep 17 00:00:00 2001 From: Michael <35783820+mib1185@users.noreply.github.com> Date: Sat, 20 Apr 2024 12:13:56 +0200 Subject: [PATCH] Automatic cleanup of entity and device registry in AVM FRITZ!SmartHome (#114601) --- homeassistant/components/fritzbox/__init__.py | 10 +- .../components/fritzbox/coordinator.py | 57 +++++++-- tests/components/fritzbox/test_coordinator.py | 111 ++++++++++++++++++ tests/components/fritzbox/test_init.py | 62 +--------- 4 files changed, 165 insertions(+), 75 deletions(-) create mode 100644 tests/components/fritzbox/test_coordinator.py diff --git a/homeassistant/components/fritzbox/__init__.py b/homeassistant/components/fritzbox/__init__.py index 7f4006768c4..904a86d21ae 100644 --- a/homeassistant/components/fritzbox/__init__.py +++ b/homeassistant/components/fritzbox/__init__.py @@ -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: diff --git a/homeassistant/components/fritzbox/coordinator.py b/homeassistant/components/fritzbox/coordinator.py index c58665f2b5d..a9cfc25b223 100644 --- a/homeassistant/components/fritzbox/coordinator.py +++ b/homeassistant/components/fritzbox/coordinator.py @@ -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: diff --git a/tests/components/fritzbox/test_coordinator.py b/tests/components/fritzbox/test_coordinator.py new file mode 100644 index 00000000000..401fab8f169 --- /dev/null +++ b/tests/components/fritzbox/test_coordinator.py @@ -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 diff --git a/tests/components/fritzbox/test_init.py b/tests/components/fritzbox/test_init.py index 4ee351f7914..8d7e4249fbd 100644 --- a/tests/components/fritzbox/test_init.py +++ b/tests/components/fritzbox/test_init.py @@ -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()