Improve entity registry handling of device changes (home-assistant#148425)

Author: emontnemery

homeassistant/helpers/device_registry.py

@@ -144,13 +144,21 @@ class DeviceInfo(TypedDict, total=False):
 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 @@ def async_remove_device(self, device_id: str) -> None:
                 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()

homeassistant/helpers/entity_registry.py

@@ -1103,8 +1103,17 @@ def async_device_modified(
             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 @@ def async_device_modified(
 
         # 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:

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

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(