Allow deleting entity entries from entity_registry.async_migrate_entries (#101094)

* Allow deleting entity entries from entity_registry.async_migrate_entries

* Explicitly return None in tests

homeassistant/helpers/entity_registry.py

@@ -1323,12 +1323,18 @@ async def async_migrate_entries(
     config_entry_id: str,
     entry_callback: Callable[[RegistryEntry], dict[str, Any] | None],
 ) -> None:
-    """Migrator of unique IDs."""
+    """Migrate entity registry entries which belong to a config entry.
+
+    Can be used as a migrator of unique_ids or to update other entity registry data.
+    Can also be used to remove duplicated entity registry entries.
+    """
     ent_reg = async_get(hass)
 
-    for entry in ent_reg.entities.values():
+    for entry in list(ent_reg.entities.values()):
         if entry.config_entry_id != config_entry_id:
             continue
+        if not ent_reg.entities.get_entry(entry.id):
+            continue
 
         updates = entry_callback(entry)
 

tests/helpers/test_entity_registry.py

@@ -1684,3 +1684,69 @@ async def test_restore_entity(hass, update_events, freezer):
     assert update_events[11] == {"action": "remove", "entity_id": "light.hue_1234"}
     # Restore entities the 3rd time
     assert update_events[12] == {"action": "create", "entity_id": "light.hue_1234"}
+
+
+async def test_async_migrate_entry_delete_self(hass):
+    """Test async_migrate_entry."""
+    registry = er.async_get(hass)
+    config_entry1 = MockConfigEntry(domain="test1")
+    config_entry2 = MockConfigEntry(domain="test2")
+    entry1 = registry.async_get_or_create(
+        "light", "hue", "1234", config_entry=config_entry1, original_name="Entry 1"
+    )
+    entry2 = registry.async_get_or_create(
+        "light", "hue", "5678", config_entry=config_entry1, original_name="Entry 2"
+    )
+    entry3 = registry.async_get_or_create(
+        "light", "hue", "90AB", config_entry=config_entry2, original_name="Entry 3"
+    )
+
+    @callback
+    def _async_migrator(entity_entry: er.RegistryEntry) -> dict[str, Any] | None:
+        entries.add(entity_entry.entity_id)
+        if entity_entry == entry1:
+            registry.async_remove(entry1.entity_id)
+            return None
+        if entity_entry == entry2:
+            return {"original_name": "Entry 2 renamed"}
+        return None
+
+    entries = set()
+    await er.async_migrate_entries(hass, config_entry1.entry_id, _async_migrator)
+    assert entries == {entry1.entity_id, entry2.entity_id}
+    assert not registry.async_is_registered(entry1.entity_id)
+    entry2 = registry.async_get(entry2.entity_id)
+    assert entry2.original_name == "Entry 2 renamed"
+    assert registry.async_get(entry3.entity_id) is entry3
+
+
+async def test_async_migrate_entry_delete_other(hass):
+    """Test async_migrate_entry."""
+    registry = er.async_get(hass)
+    config_entry1 = MockConfigEntry(domain="test1")
+    config_entry2 = MockConfigEntry(domain="test2")
+    entry1 = registry.async_get_or_create(
+        "light", "hue", "1234", config_entry=config_entry1, original_name="Entry 1"
+    )
+    entry2 = registry.async_get_or_create(
+        "light", "hue", "5678", config_entry=config_entry1, original_name="Entry 2"
+    )
+    registry.async_get_or_create(
+        "light", "hue", "90AB", config_entry=config_entry2, original_name="Entry 3"
+    )
+
+    @callback
+    def _async_migrator(entity_entry: er.RegistryEntry) -> dict[str, Any] | None:
+        entries.add(entity_entry.entity_id)
+        if entity_entry == entry1:
+            registry.async_remove(entry2.entity_id)
+            return None
+        if entity_entry == entry2:
+            # We should not get here
+            pytest.fail()
+        return None
+
+    entries = set()
+    await er.async_migrate_entries(hass, config_entry1.entry_id, _async_migrator)
+    assert entries == {entry1.entity_id}
+    assert not registry.async_is_registered(entry2.entity_id)