From 7a6facc875ac0f3af33450d77880777ffcb3816d Mon Sep 17 00:00:00 2001
From: Robert Svensson <Kane610@users.noreply.github.com>
Date: Tue, 4 Sep 2018 09:00:14 +0200
Subject: [PATCH] Device and entity registry remove config entry on unload
(#16247)
* Test
* Ability to remove device
* Don't remove devices, instead remove config entry from device and entity registries
* Remove print
* Remove is not the same as unload
* Add tests
* Fix hound comment
---
homeassistant/config_entries.py | 8 +++++++
homeassistant/helpers/device_registry.py | 10 ++++++++-
homeassistant/helpers/entity_registry.py | 10 ++++++++-
tests/helpers/test_device_registry.py | 27 ++++++++++++++++++++++++
tests/helpers/test_entity_registry.py | 9 ++++++++
5 files changed, 62 insertions(+), 2 deletions(-)
diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py
index 8db09cdb8da..6eae9e13030 100644
--- a/homeassistant/config_entries.py
+++ b/homeassistant/config_entries.py
@@ -325,6 +325,14 @@ class ConfigEntries:
unloaded = await entry.async_unload(self.hass)
+ device_registry = await \
+ self.hass.helpers.device_registry.async_get_registry()
+ device_registry.async_clear_config_entry(entry_id)
+
+ entity_registry = await \
+ self.hass.helpers.entity_registry.async_get_registry()
+ entity_registry.async_clear_config_entry(entry_id)
+
return {
'require_restart': not unloaded
}
diff --git a/homeassistant/helpers/device_registry.py b/homeassistant/helpers/device_registry.py
index 504448b948d..e6ff45af2fe 100644
--- a/homeassistant/helpers/device_registry.py
+++ b/homeassistant/helpers/device_registry.py
@@ -45,7 +45,7 @@ class DeviceRegistry:
self._store = hass.helpers.storage.Store(STORAGE_VERSION, STORAGE_KEY)
@callback
- def async_get_device(self, identifiers: str, connections: tuple):
+ def async_get_device(self, identifiers: set, connections: set):
"""Check if device is registered."""
for device in self.devices.values():
if any(iden in device.identifiers for iden in identifiers) or \
@@ -127,6 +127,14 @@ class DeviceRegistry:
return data
+ @callback
+ def async_clear_config_entry(self, config_entry):
+ """Clear config entry from registry entries."""
+ for device in self.devices.values():
+ if config_entry in device.config_entries:
+ device.config_entries.remove(config_entry)
+ self.async_schedule_save()
+
@bind_hass
async def async_get_registry(hass) -> DeviceRegistry:
diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py
index 804ee4235d0..da3645a96fe 100644
--- a/homeassistant/helpers/entity_registry.py
+++ b/homeassistant/helpers/entity_registry.py
@@ -31,7 +31,7 @@ STORAGE_VERSION = 1
STORAGE_KEY = 'core.entity_registry'
-@attr.s(slots=True, frozen=True)
+@attr.s(slots=True)
class RegistryEntry:
"""Entity Registry Entry."""
@@ -250,6 +250,14 @@ class EntityRegistry:
return data
+ @callback
+ def async_clear_config_entry(self, config_entry):
+ """Clear config entry from registry entries."""
+ for entry in self.entities.values():
+ if config_entry == entry.config_entry_id:
+ entry.config_entry_id = None
+ self.async_schedule_save()
+
@bind_hass
async def async_get_registry(hass) -> EntityRegistry:
diff --git a/tests/helpers/test_device_registry.py b/tests/helpers/test_device_registry.py
index 84ad54f7b82..a9132529bc3 100644
--- a/tests/helpers/test_device_registry.py
+++ b/tests/helpers/test_device_registry.py
@@ -138,3 +138,30 @@ async def test_loading_from_storage(hass, hass_storage):
manufacturer='manufacturer', model='model')
assert entry.id == 'abcdefghijklm'
assert isinstance(entry.config_entries, set)
+
+
+async def test_removing_config_entries(registry):
+ """Make sure we do not get duplicate entries."""
+ entry = registry.async_get_or_create(
+ config_entry='123',
+ connections={('ethernet', '12:34:56:78:90:AB:CD:EF')},
+ identifiers={('bridgeid', '0123')},
+ manufacturer='manufacturer', model='model')
+ entry2 = registry.async_get_or_create(
+ config_entry='456',
+ connections={('ethernet', '12:34:56:78:90:AB:CD:EF')},
+ identifiers={('bridgeid', '0123')},
+ manufacturer='manufacturer', model='model')
+ entry3 = registry.async_get_or_create(
+ config_entry='123',
+ connections={('ethernet', '34:56:78:90:AB:CD:EF:12')},
+ identifiers={('bridgeid', '4567')},
+ manufacturer='manufacturer', model='model')
+
+ assert len(registry.devices) == 2
+ assert entry is entry2
+ assert entry is not entry3
+ assert entry.config_entries == {'123', '456'}
+ registry.async_clear_config_entry('123')
+ assert entry.config_entries == {'456'}
+ assert entry3.config_entries == set()
diff --git a/tests/helpers/test_entity_registry.py b/tests/helpers/test_entity_registry.py
index d0c088a6f69..bb28287ddd8 100644
--- a/tests/helpers/test_entity_registry.py
+++ b/tests/helpers/test_entity_registry.py
@@ -186,6 +186,15 @@ async def test_updating_config_entry_id(registry):
assert entry2.config_entry_id == 'mock-id-2'
+async def test_removing_config_entry_id(registry):
+ """Test that we update config entry id in registry."""
+ entry = registry.async_get_or_create(
+ 'light', 'hue', '5678', config_entry_id='mock-id-1')
+ assert entry.config_entry_id == 'mock-id-1'
+ registry.async_clear_config_entry('mock-id-1')
+ assert entry.config_entry_id is None
+
+
async def test_migration(hass):
"""Test migration from old data to new."""
old_conf = {
--
GitLab