Automatic cleanup of entity and device registry in AVM FRITZ!SmartHome

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 @@ def _update_unique_id(entry: RegistryEntry) -> dict[str, str] | None:
 
     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:

homeassistant/components/fritzbox/coordinator.py

@@ -12,6 +12,7 @@
 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,28 +29,56 @@ 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]
-        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,
+            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()
+
         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)
+        )
+
+    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."""
         try:
@@ -95,6 +124,12 @@ def _update_fritz_devices(self) -> FritzboxCoordinatorData:
         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:

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

tests/components/fritzbox/test_init.py

@@ -6,7 +6,7 @@
 
 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()