From 2dca78efbb403e7a48363017a21d915c206b9648 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Fri, 11 Jul 2025 20:56:50 +0200 Subject: [PATCH] Improve entity registry handling of device changes (#148425) --- homeassistant/helpers/device_registry.py | 21 ++++++--- homeassistant/helpers/entity_registry.py | 60 +++++++++++++++--------- tests/helpers/test_device_registry.py | 13 ++++- tests/helpers/test_entity_registry.py | 35 ++++++++------ 4 files changed, 87 insertions(+), 42 deletions(-) diff --git a/homeassistant/helpers/device_registry.py b/homeassistant/helpers/device_registry.py index bad772abaff..bc6e7c810bf 100644 --- a/homeassistant/helpers/device_registry.py +++ b/homeassistant/helpers/device_registry.py @@ -144,13 +144,21 @@ DEVICE_INFO_KEYS = set.union(*(itm for itm in DEVICE_INFO_TYPES.values())) LOW_PRIO_CONFIG_ENTRY_DOMAINS = {"homekit_controller", "matter", "mqtt", "upnp"} -class _EventDeviceRegistryUpdatedData_CreateRemove(TypedDict): - """EventDeviceRegistryUpdated data for action type 'create' and 'remove'.""" +class _EventDeviceRegistryUpdatedData_Create(TypedDict): + """EventDeviceRegistryUpdated data for action type 'create'.""" - action: Literal["create", "remove"] + action: Literal["create"] device_id: str +class _EventDeviceRegistryUpdatedData_Remove(TypedDict): + """EventDeviceRegistryUpdated data for action type 'remove'.""" + + action: Literal["remove"] + device_id: str + device: DeviceEntry + + class _EventDeviceRegistryUpdatedData_Update(TypedDict): """EventDeviceRegistryUpdated data for action type 'update'.""" @@ -160,7 +168,8 @@ class _EventDeviceRegistryUpdatedData_Update(TypedDict): type EventDeviceRegistryUpdatedData = ( - _EventDeviceRegistryUpdatedData_CreateRemove + _EventDeviceRegistryUpdatedData_Create + | _EventDeviceRegistryUpdatedData_Remove | _EventDeviceRegistryUpdatedData_Update ) @@ -1309,8 +1318,8 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]): self.async_update_device(other_device.id, via_device_id=None) self.hass.bus.async_fire_internal( EVENT_DEVICE_REGISTRY_UPDATED, - _EventDeviceRegistryUpdatedData_CreateRemove( - action="remove", device_id=device_id + _EventDeviceRegistryUpdatedData_Remove( + action="remove", device_id=device_id, device=device ), ) self.async_schedule_save() diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index 0b61c3e8f16..ddb25c7b0a8 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -1103,8 +1103,17 @@ class EntityRegistry(BaseRegistry): entities = async_entries_for_device( self, event.data["device_id"], include_disabled_entities=True ) + removed_device = event.data["device"] for entity in entities: - self.async_remove(entity.entity_id) + config_entry_id = entity.config_entry_id + if ( + config_entry_id in removed_device.config_entries + and entity.config_subentry_id + in removed_device.config_entries_subentries[config_entry_id] + ): + self.async_remove(entity.entity_id) + else: + self.async_update_entity(entity.entity_id, device_id=None) return if event.data["action"] != "update": @@ -1121,29 +1130,38 @@ class EntityRegistry(BaseRegistry): # Remove entities which belong to config entries no longer associated with the # device - entities = async_entries_for_device( - self, event.data["device_id"], include_disabled_entities=True - ) - for entity in entities: - if ( - entity.config_entry_id is not None - and entity.config_entry_id not in device.config_entries - ): - self.async_remove(entity.entity_id) + if old_config_entries := event.data["changes"].get("config_entries"): + entities = async_entries_for_device( + self, event.data["device_id"], include_disabled_entities=True + ) + for entity in entities: + config_entry_id = entity.config_entry_id + if ( + entity.config_entry_id in old_config_entries + and entity.config_entry_id not in device.config_entries + ): + self.async_remove(entity.entity_id) # Remove entities which belong to config subentries no longer associated with the # device - entities = async_entries_for_device( - self, event.data["device_id"], include_disabled_entities=True - ) - for entity in entities: - if ( - (config_entry_id := entity.config_entry_id) is not None - and config_entry_id in device.config_entries - and entity.config_subentry_id - not in device.config_entries_subentries[config_entry_id] - ): - self.async_remove(entity.entity_id) + if old_config_entries_subentries := event.data["changes"].get( + "config_entries_subentries" + ): + entities = async_entries_for_device( + self, event.data["device_id"], include_disabled_entities=True + ) + for entity in entities: + config_entry_id = entity.config_entry_id + config_subentry_id = entity.config_subentry_id + if ( + config_entry_id in device.config_entries + and config_entry_id in old_config_entries_subentries + and config_subentry_id + in old_config_entries_subentries[config_entry_id] + and config_subentry_id + not in device.config_entries_subentries[config_entry_id] + ): + self.async_remove(entity.entity_id) # Re-enable disabled entities if the device is no longer disabled if not device.disabled: diff --git a/tests/helpers/test_device_registry.py b/tests/helpers/test_device_registry.py index 58933ca4314..23a451dd06c 100644 --- a/tests/helpers/test_device_registry.py +++ b/tests/helpers/test_device_registry.py @@ -1652,6 +1652,7 @@ async def test_removing_config_entries( assert update_events[4].data == { "action": "remove", "device_id": entry3.id, + "device": entry3, } @@ -1724,10 +1725,12 @@ async def test_deleted_device_removing_config_entries( assert update_events[3].data == { "action": "remove", "device_id": entry.id, + "device": entry2, } assert update_events[4].data == { "action": "remove", "device_id": entry3.id, + "device": entry3, } device_registry.async_clear_config_entry(config_entry_1.entry_id) @@ -1973,6 +1976,7 @@ async def test_removing_config_subentries( assert update_events[7].data == { "action": "remove", "device_id": entry.id, + "device": entry, } @@ -2102,6 +2106,7 @@ async def test_deleted_device_removing_config_subentries( assert update_events[4].data == { "action": "remove", "device_id": entry.id, + "device": entry4, } device_registry.async_clear_config_subentry(config_entry_1.entry_id, None) @@ -2925,6 +2930,7 @@ async def test_update_remove_config_entries( assert update_events[6].data == { "action": "remove", "device_id": entry3.id, + "device": entry3, } @@ -3104,6 +3110,7 @@ async def test_update_remove_config_subentries( config_entry_3.entry_id: {None}, } + entry_before_remove = entry entry = device_registry.async_update_device( entry_id, remove_config_entry_id=config_entry_3.entry_id, @@ -3201,6 +3208,7 @@ async def test_update_remove_config_subentries( assert update_events[7].data == { "action": "remove", "device_id": entry_id, + "device": entry_before_remove, } @@ -3422,7 +3430,7 @@ async def test_restore_device( ) # Apply user customizations - device_registry.async_update_device( + entry = device_registry.async_update_device( entry.id, area_id="12345A", disabled_by=dr.DeviceEntryDisabler.USER, @@ -3543,6 +3551,7 @@ async def test_restore_device( assert update_events[2].data == { "action": "remove", "device_id": entry.id, + "device": entry, } assert update_events[3].data == { "action": "create", @@ -3865,6 +3874,7 @@ async def test_restore_shared_device( assert update_events[3].data == { "action": "remove", "device_id": entry.id, + "device": updated_device, } assert update_events[4].data == { "action": "create", @@ -3873,6 +3883,7 @@ async def test_restore_shared_device( assert update_events[5].data == { "action": "remove", "device_id": entry.id, + "device": entry2, } assert update_events[6].data == { "action": "create", diff --git a/tests/helpers/test_entity_registry.py b/tests/helpers/test_entity_registry.py index 5afffebb5f6..40a26295cbb 100644 --- a/tests/helpers/test_entity_registry.py +++ b/tests/helpers/test_entity_registry.py @@ -1684,20 +1684,23 @@ async def test_remove_config_entry_from_device_removes_entities_2( await hass.async_block_till_done() assert device_registry.async_get(device_entry.id) + # Entities which are not tied to the removed config entry should not be removed assert entity_registry.async_is_registered(entry_1.entity_id) - # Entities with a config entry not in the device are removed - assert not entity_registry.async_is_registered(entry_2.entity_id) + assert entity_registry.async_is_registered(entry_2.entity_id) - # Remove the second config entry from the device + # Remove the second config entry from the device (this removes the device) device_registry.async_update_device( device_entry.id, remove_config_entry_id=config_entry_2.entry_id ) await hass.async_block_till_done() assert not device_registry.async_get(device_entry.id) - # The device is removed, both entities are now removed - assert not entity_registry.async_is_registered(entry_1.entity_id) - assert not entity_registry.async_is_registered(entry_2.entity_id) + # Entities which are not tied to a config entry in the device should not be removed + assert entity_registry.async_is_registered(entry_1.entity_id) + assert entity_registry.async_is_registered(entry_2.entity_id) + # Check the device link is set to None + assert entity_registry.async_get(entry_1.entity_id).device_id is None + assert entity_registry.async_get(entry_2.entity_id).device_id is None async def test_remove_config_subentry_from_device_removes_entities( @@ -1921,12 +1924,12 @@ async def test_remove_config_subentry_from_device_removes_entities_2( await hass.async_block_till_done() assert device_registry.async_get(device_entry.id) + # Entities with a config subentry not in the device are not removed assert entity_registry.async_is_registered(entry_1.entity_id) - # Entities with a config subentry not in the device are removed - assert not entity_registry.async_is_registered(entry_2.entity_id) - assert not entity_registry.async_is_registered(entry_3.entity_id) + assert entity_registry.async_is_registered(entry_2.entity_id) + assert entity_registry.async_is_registered(entry_3.entity_id) - # Remove the second config subentry from the device + # Remove the second config subentry from the device, this removes the device device_registry.async_update_device( device_entry.id, remove_config_entry_id=config_entry_1.entry_id, @@ -1935,10 +1938,14 @@ async def test_remove_config_subentry_from_device_removes_entities_2( await hass.async_block_till_done() assert not device_registry.async_get(device_entry.id) - # All entities are now removed - assert not entity_registry.async_is_registered(entry_1.entity_id) - assert not entity_registry.async_is_registered(entry_2.entity_id) - assert not entity_registry.async_is_registered(entry_3.entity_id) + # Entities with a config subentry not in the device are not removed + assert entity_registry.async_is_registered(entry_1.entity_id) + assert entity_registry.async_is_registered(entry_2.entity_id) + assert entity_registry.async_is_registered(entry_3.entity_id) + # Check the device link is set to None + assert entity_registry.async_get(entry_1.entity_id).device_id is None + assert entity_registry.async_get(entry_2.entity_id).device_id is None + assert entity_registry.async_get(entry_3.entity_id).device_id is None async def test_update_device_race(