Add entity registry helper to update entity platform (#69162)

* Add entity registry helper to migrate entity to new platform

* Add additional assertion

* Add more properties to migration logic

* Change logic after thinking about erik's comments

* Require new_config_entry_id if entry.config_entry_id is not None

* Create private async_update_entity function that all update functions use

* Don't have special handling for entity ID missing in async_update_entity_platform

* fix docstring

homeassistant/components/tomorrowio/__init__.py

@@ -155,36 +155,27 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         for entity_entry in er.async_entries_for_config_entry(
             ent_reg, old_config_entry_id
         ):
-            _LOGGER.debug("Removing %s", entity_entry.entity_id)
+            old_platform = entity_entry.platform
-            ent_reg.async_remove(entity_entry.entity_id)
             # In case the API key has changed due to a V3 -> V4 change, we need to
             # generate the new entity's unique ID
             new_unique_id = (
                 f"{entry.data[CONF_API_KEY]}_"
                 f"{'_'.join(entity_entry.unique_id.split('_')[1:])}"
             )
-            _LOGGER.debug(
+            ent_reg.async_update_entity_platform(
-                "Re-creating %s for the new config entry", entity_entry.entity_id
+                entity_entry.entity_id,
-            )
+                DOMAIN,
-            # We will precreate the entity so that any customizations can be preserved
+                new_unique_id=new_unique_id,
-            new_entity_entry = ent_reg.async_get_or_create(
+                new_config_entry_id=entry.entry_id,
-                entity_entry.domain,
+                new_device_id=device.id,
+            )
+            assert entity_entry
+            _LOGGER.debug(
+                "Migrated %s from %s to %s",
+                entity_entry.entity_id,
+                old_platform,
                 DOMAIN,
-                new_unique_id,
-                suggested_object_id=entity_entry.entity_id.split(".")[1],
-                disabled_by=entity_entry.disabled_by,
-                config_entry=entry,
-                original_name=entity_entry.original_name,
-                original_icon=entity_entry.original_icon,
             )
-            _LOGGER.debug("Re-created %s", new_entity_entry.entity_id)
-            # If there are customizations on the old entity, apply them to the new one
-            if entity_entry.name or entity_entry.icon:
-                ent_reg.async_update_entity(
-                    new_entity_entry.entity_id,
-                    name=entity_entry.name,
-                    icon=entity_entry.icon,
-                )
 
         # We only have one device in the registry but we will do a loop just in case
         for old_device in dr.async_entries_for_config_entry(

homeassistant/helpers/entity_registry.py

@@ -29,6 +29,7 @@ from homeassistant.const import (
     MAX_LENGTH_STATE_DOMAIN,
     MAX_LENGTH_STATE_ENTITY_ID,
     STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
 )
 from homeassistant.core import (
     Event,
@@ -484,7 +485,7 @@ class EntityRegistry:
             )
 
     @callback
-    def async_update_entity(
+    def _async_update_entity(
         self,
         entity_id: str,
         *,
@@ -505,6 +506,8 @@ class EntityRegistry:
         original_name: str | None | UndefinedType = UNDEFINED,
         supported_features: int | UndefinedType = UNDEFINED,
         unit_of_measurement: str | None | UndefinedType = UNDEFINED,
+        platform: str | None | UndefinedType = UNDEFINED,
+        options: Mapping[str, Mapping[str, Any]] | UndefinedType = UNDEFINED,
     ) -> RegistryEntry:
         """Private facing update properties method."""
         old = self.entities[entity_id]
@@ -544,6 +547,8 @@ class EntityRegistry:
             ("original_name", original_name),
             ("supported_features", supported_features),
             ("unit_of_measurement", unit_of_measurement),
+            ("platform", platform),
+            ("options", options),
         ):
             if value is not UNDEFINED and value != getattr(old, attr_name):
                 new_values[attr_name] = value
@@ -595,6 +600,87 @@ class EntityRegistry:
 
         return new
 
+    @callback
+    def async_update_entity(
+        self,
+        entity_id: str,
+        *,
+        area_id: str | None | UndefinedType = UNDEFINED,
+        capabilities: Mapping[str, Any] | None | UndefinedType = UNDEFINED,
+        config_entry_id: str | None | UndefinedType = UNDEFINED,
+        device_class: str | None | UndefinedType = UNDEFINED,
+        device_id: str | None | UndefinedType = UNDEFINED,
+        disabled_by: RegistryEntryDisabler | None | UndefinedType = UNDEFINED,
+        entity_category: EntityCategory | None | UndefinedType = UNDEFINED,
+        hidden_by: RegistryEntryHider | None | UndefinedType = UNDEFINED,
+        icon: str | None | UndefinedType = UNDEFINED,
+        name: str | None | UndefinedType = UNDEFINED,
+        new_entity_id: str | UndefinedType = UNDEFINED,
+        new_unique_id: str | UndefinedType = UNDEFINED,
+        original_device_class: str | None | UndefinedType = UNDEFINED,
+        original_icon: str | None | UndefinedType = UNDEFINED,
+        original_name: str | None | UndefinedType = UNDEFINED,
+        supported_features: int | UndefinedType = UNDEFINED,
+        unit_of_measurement: str | None | UndefinedType = UNDEFINED,
+    ) -> RegistryEntry:
+        """Update properties of an entity."""
+        return self._async_update_entity(
+            entity_id,
+            area_id=area_id,
+            capabilities=capabilities,
+            config_entry_id=config_entry_id,
+            device_class=device_class,
+            device_id=device_id,
+            disabled_by=disabled_by,
+            entity_category=entity_category,
+            hidden_by=hidden_by,
+            icon=icon,
+            name=name,
+            new_entity_id=new_entity_id,
+            new_unique_id=new_unique_id,
+            original_device_class=original_device_class,
+            original_icon=original_icon,
+            original_name=original_name,
+            supported_features=supported_features,
+            unit_of_measurement=unit_of_measurement,
+        )
+
+    @callback
+    def async_update_entity_platform(
+        self,
+        entity_id: str,
+        new_platform: str,
+        *,
+        new_config_entry_id: str | UndefinedType = UNDEFINED,
+        new_unique_id: str | UndefinedType = UNDEFINED,
+        new_device_id: str | None | UndefinedType = UNDEFINED,
+    ) -> RegistryEntry:
+        """
+        Update entity platform.
+
+        This should only be used when an entity needs to be migrated between
+        integrations.
+        """
+        if (
+            state := self.hass.states.get(entity_id)
+        ) is not None and state.state != STATE_UNKNOWN:
+            raise ValueError("Only entities that haven't been loaded can be migrated")
+
+        old = self.entities[entity_id]
+        if new_config_entry_id == UNDEFINED and old.config_entry_id is not None:
+            raise ValueError(
+                f"new_config_entry_id required because {entity_id} is already linked "
+                "to a config entry"
+            )
+
+        return self._async_update_entity(
+            entity_id,
+            new_unique_id=new_unique_id,
+            config_entry_id=new_config_entry_id,
+            device_id=new_device_id,
+            platform=new_platform,
+        )
+
     @callback
     def async_update_entity_options(
         self, entity_id: str, domain: str, options: dict[str, Any]
@@ -602,19 +688,7 @@ class EntityRegistry:
         """Update entity options."""
         old = self.entities[entity_id]
         new_options: Mapping[str, Mapping[str, Any]] = {**old.options, domain: options}
-        new = self.entities[entity_id] = attr.evolve(old, options=new_options)
+        return self._async_update_entity(entity_id, options=new_options)
-
-        self.async_schedule_save()
-
-        data: dict[str, str | dict[str, Any]] = {
-            "action": "update",
-            "entity_id": entity_id,
-            "changes": {"options": old.options},
-        }
-
-        self.hass.bus.async_fire(EVENT_ENTITY_REGISTRY_UPDATED, data)
-
-        return new
 
     async def async_load(self) -> None:
         """Load the entity registry."""

tests/helpers/test_entity_registry.py

@@ -1247,3 +1247,74 @@ async def test_entity_category_str_not_allowed(hass):
         reg.async_update_entity(
             entity_id, entity_category=EntityCategory.DIAGNOSTIC.value
         )
+
+
+def test_migrate_entity_to_new_platform(hass, registry):
+    """Test migrate_entity_to_new_platform."""
+    orig_config_entry = MockConfigEntry(domain="light")
+    orig_unique_id = "5678"
+
+    orig_entry = registry.async_get_or_create(
+        "light",
+        "hue",
+        orig_unique_id,
+        suggested_object_id="light",
+        config_entry=orig_config_entry,
+        disabled_by=er.RegistryEntryDisabler.USER,
+        entity_category=EntityCategory.CONFIG,
+        original_device_class="mock-device-class",
+        original_icon="initial-original_icon",
+        original_name="initial-original_name",
+    )
+    assert registry.async_get("light.light") is orig_entry
+    registry.async_update_entity(
+        "light.light",
+        name="new_name",
+        icon="new_icon",
+    )
+
+    new_config_entry = MockConfigEntry(domain="light")
+    new_unique_id = "1234"
+
+    assert registry.async_update_entity_platform(
+        "light.light",
+        "hue2",
+        new_unique_id=new_unique_id,
+        new_config_entry_id=new_config_entry.entry_id,
+    )
+
+    assert not registry.async_get_entity_id("light", "hue", orig_unique_id)
+
+    assert (new_entry := registry.async_get("light.light")) is not orig_entry
+
+    assert new_entry.config_entry_id == new_config_entry.entry_id
+    assert new_entry.unique_id == new_unique_id
+    assert new_entry.name == "new_name"
+    assert new_entry.icon == "new_icon"
+    assert new_entry.platform == "hue2"
+
+    # Test nonexisting entity
+    with pytest.raises(KeyError):
+        registry.async_update_entity_platform(
+            "light.not_a_real_light",
+            "hue2",
+            new_unique_id=new_unique_id,
+            new_config_entry_id=new_config_entry.entry_id,
+        )
+
+    # Test migrate entity without new config entry ID
+    with pytest.raises(ValueError):
+        registry.async_update_entity_platform(
+            "light.light",
+            "hue3",
+        )
+
+    # Test entity with a state
+    hass.states.async_set("light.light", "on")
+    with pytest.raises(ValueError):
+        registry.async_update_entity_platform(
+            "light.light",
+            "hue2",
+            new_unique_id=new_unique_id,
+            new_config_entry_id=new_config_entry.entry_id,
+        )