From 4ae31bc9384c850a624ec721890672ef20d3226e Mon Sep 17 00:00:00 2001 From: Paulus Schoutsen Date: Tue, 5 May 2020 10:53:46 -0700 Subject: [PATCH] Clean up device registry if entity registry updates (#35106) --- homeassistant/helpers/debounce.py | 2 +- homeassistant/helpers/device_registry.py | 90 ++++++++++++++++++++---- tests/helpers/test_device_registry.py | 83 ++++++++++++++++++++-- 3 files changed, 157 insertions(+), 18 deletions(-) diff --git a/homeassistant/helpers/debounce.py b/homeassistant/helpers/debounce.py index 6206081dc8c..3f297dcbbe8 100644 --- a/homeassistant/helpers/debounce.py +++ b/homeassistant/helpers/debounce.py @@ -21,7 +21,7 @@ class Debouncer: """Initialize debounce. immediate: indicate if the function needs to be called right away and - wait 0.3s until executing next invocation. + wait until executing next invocation. function: optional and can be instantiated later. """ self.hass = hass diff --git a/homeassistant/helpers/device_registry.py b/homeassistant/helpers/device_registry.py index 7b49f6bd363..8fbb81962ff 100644 --- a/homeassistant/helpers/device_registry.py +++ b/homeassistant/helpers/device_registry.py @@ -1,16 +1,21 @@ """Provide a way to connect entities belonging to one device.""" from collections import OrderedDict import logging -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple import uuid import attr -from homeassistant.core import callback +from homeassistant.const import EVENT_HOMEASSISTANT_STARTED +from homeassistant.core import Event, callback +from .debounce import Debouncer from .singleton import singleton from .typing import HomeAssistantType +if TYPE_CHECKING: + from . import entity_registry + # mypy: allow-untyped-calls, allow-untyped-defs, no-check-untyped-defs _LOGGER = logging.getLogger(__name__) @@ -21,6 +26,7 @@ EVENT_DEVICE_REGISTRY_UPDATED = "device_registry_updated" STORAGE_KEY = "core.device_registry" STORAGE_VERSION = 1 SAVE_DELAY = 10 +CLEANUP_DELAY = 10 CONNECTION_NETWORK_MAC = "mac" CONNECTION_UPNP = "upnp" @@ -285,6 +291,8 @@ class DeviceRegistry: async def async_load(self): """Load the device registry.""" + async_setup_cleanup(self.hass, self) + data = await self._store.async_load() devices = OrderedDict() @@ -347,16 +355,8 @@ class DeviceRegistry: @callback def async_clear_config_entry(self, config_entry_id: str) -> None: """Clear config entry from registry entries.""" - remove = [] - for dev_id, device in self.devices.items(): - if device.config_entries == {config_entry_id}: - remove.append(dev_id) - else: - self._async_update_device( - dev_id, remove_config_entry_id=config_entry_id - ) - for dev_id in remove: - self.async_remove_device(dev_id) + for device in list(self.devices.values()): + self._async_update_device(device.id, remove_config_entry_id=config_entry_id) @callback def async_clear_area_id(self, area_id: str) -> None: @@ -390,3 +390,69 @@ def async_entries_for_config_entry( for device in registry.devices.values() if config_entry_id in device.config_entries ] + + +@callback +def async_cleanup( + hass: HomeAssistantType, + dev_reg: DeviceRegistry, + ent_reg: "entity_registry.EntityRegistry", +) -> None: + """Clean up device registry.""" + # Find all devices that are no longer referenced in the entity registry. + referenced = {entry.device_id for entry in ent_reg.entities.values()} + orphan = set(dev_reg.devices) - referenced + + for dev_id in orphan: + dev_reg.async_remove_device(dev_id) + + # Find all referenced config entries that no longer exist + # This shouldn't happen but have not been able to track down the bug :( + config_entry_ids = {entry.entry_id for entry in hass.config_entries.async_entries()} + + for device in list(dev_reg.devices.values()): + for config_entry_id in device.config_entries: + if config_entry_id not in config_entry_ids: + dev_reg.async_update_device( + device.id, remove_config_entry_id=config_entry_id + ) + + +@callback +def async_setup_cleanup(hass: HomeAssistantType, dev_reg: DeviceRegistry) -> None: + """Clean up device registry when entities removed.""" + from . import entity_registry # pylint: disable=import-outside-toplevel + + async def cleanup(): + """Cleanup.""" + ent_reg = await entity_registry.async_get_registry(hass) + async_cleanup(hass, dev_reg, ent_reg) + + debounced_cleanup = Debouncer( + hass, _LOGGER, cooldown=CLEANUP_DELAY, immediate=False, function=cleanup + ) + + async def entity_registry_changed(event: Event) -> None: + """Handle entity updated or removed.""" + if ( + event.data["action"] == "update" + and "device_id" not in event.data["changes"] + ) or event.data["action"] == "create": + return + + await debounced_cleanup.async_call() + + if hass.is_running: + hass.bus.async_listen( + entity_registry.EVENT_ENTITY_REGISTRY_UPDATED, entity_registry_changed + ) + return + + async def startup_clean(event: Event) -> None: + """Clean up on startup.""" + hass.bus.async_listen( + entity_registry.EVENT_ENTITY_REGISTRY_UPDATED, entity_registry_changed + ) + await debounced_cleanup.async_call() + + hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STARTED, startup_clean) diff --git a/tests/helpers/test_device_registry.py b/tests/helpers/test_device_registry.py index 2823ea6ca07..ef5f92de79c 100644 --- a/tests/helpers/test_device_registry.py +++ b/tests/helpers/test_device_registry.py @@ -3,12 +3,12 @@ import asyncio import pytest -from homeassistant.core import callback -from homeassistant.helpers import device_registry +from homeassistant.const import EVENT_HOMEASSISTANT_STARTED +from homeassistant.core import CoreState, callback +from homeassistant.helpers import device_registry, entity_registry -import tests.async_mock from tests.async_mock import patch -from tests.common import flush_store, mock_device_registry +from tests.common import MockConfigEntry, flush_store, mock_device_registry @pytest.fixture @@ -483,7 +483,7 @@ async def test_update_remove_config_entries(hass, registry, update_events): async def test_loading_race_condition(hass): """Test only one storage load called when concurrent loading occurred .""" - with tests.async_mock.patch( + with patch( "homeassistant.helpers.device_registry.DeviceRegistry.async_load" ) as mock_load: results = await asyncio.gather( @@ -511,3 +511,76 @@ async def test_update_sw_version(registry): assert mock_save.call_count == 1 assert updated_entry != entry assert updated_entry.sw_version == sw_version + + +async def test_cleanup_device_registry(hass, registry): + """Test cleanup works.""" + config_entry = MockConfigEntry(domain="hue") + config_entry.add_to_hass(hass) + + d1 = registry.async_get_or_create( + identifiers={("hue", "d1")}, config_entry_id=config_entry.entry_id + ) + registry.async_get_or_create( + identifiers={("hue", "d2")}, config_entry_id=config_entry.entry_id + ) + d3 = registry.async_get_or_create( + identifiers={("hue", "d3")}, config_entry_id=config_entry.entry_id + ) + registry.async_get_or_create( + identifiers={("something", "d4")}, config_entry_id="non_existing" + ) + + ent_reg = await entity_registry.async_get_registry(hass) + ent_reg.async_get_or_create("light", "hue", "e1", device_id=d1.id) + ent_reg.async_get_or_create("light", "hue", "e2", device_id=d1.id) + ent_reg.async_get_or_create("light", "hue", "e3", device_id=d3.id) + + device_registry.async_cleanup(hass, registry, ent_reg) + + assert registry.async_get_device({("hue", "d1")}, set()) is not None + assert registry.async_get_device({("hue", "d2")}, set()) is None + assert registry.async_get_device({("hue", "d3")}, set()) is not None + assert registry.async_get_device({("something", "d4")}, set()) is None + + +async def test_cleanup_startup(hass): + """Test we run a cleanup on startup.""" + hass.state = CoreState.not_running + await device_registry.async_get_registry(hass) + + with patch( + "homeassistant.helpers.device_registry.Debouncer.async_call" + ) as mock_call: + hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED) + await hass.async_block_till_done() + + assert len(mock_call.mock_calls) == 1 + + +async def test_cleanup_entity_registry_change(hass): + """Test we run a cleanup when entity registry changes.""" + await device_registry.async_get_registry(hass) + ent_reg = await entity_registry.async_get_registry(hass) + + with patch( + "homeassistant.helpers.device_registry.Debouncer.async_call" + ) as mock_call: + entity = ent_reg.async_get_or_create("light", "hue", "e1") + await hass.async_block_till_done() + assert len(mock_call.mock_calls) == 0 + + # Normal update does not trigger + ent_reg.async_update_entity(entity.entity_id, name="updated") + await hass.async_block_till_done() + assert len(mock_call.mock_calls) == 0 + + # Device ID update triggers + ent_reg.async_get_or_create("light", "hue", "e1", device_id="bla") + await hass.async_block_till_done() + assert len(mock_call.mock_calls) == 1 + + # Removal also triggers + ent_reg.async_remove(entity.entity_id) + await hass.async_block_till_done() + assert len(mock_call.mock_calls) == 2