diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index 42de4749215..bd3077c1d59 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/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) diff --git a/tests/helpers/test_entity_registry.py b/tests/helpers/test_entity_registry.py index f62addb9a64..4bf03b4d39b 100644 --- a/tests/helpers/test_entity_registry.py +++ b/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)