Improve entity registry handling of device changes (#148425)

This commit is contained in:
Erik Montnemery 2025-07-11 20:56:50 +02:00 committed by GitHub
parent e0179a7d45
commit 2dca78efbb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 87 additions and 42 deletions

View File

@ -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()

View File

@ -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:

View File

@ -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",

View File

@ -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(