mirror of
https://github.com/home-assistant/core.git
synced 2025-07-17 02:07:09 +00:00
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
This commit is contained in:
parent
42e0bc849c
commit
3bcd921a28
@ -155,36 +155,27 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
|||||||
for entity_entry in er.async_entries_for_config_entry(
|
for entity_entry in er.async_entries_for_config_entry(
|
||||||
ent_reg, old_config_entry_id
|
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
|
# In case the API key has changed due to a V3 -> V4 change, we need to
|
||||||
# generate the new entity's unique ID
|
# generate the new entity's unique ID
|
||||||
new_unique_id = (
|
new_unique_id = (
|
||||||
f"{entry.data[CONF_API_KEY]}_"
|
f"{entry.data[CONF_API_KEY]}_"
|
||||||
f"{'_'.join(entity_entry.unique_id.split('_')[1:])}"
|
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,
|
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
|
# 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(
|
for old_device in dr.async_entries_for_config_entry(
|
||||||
|
@ -29,6 +29,7 @@ from homeassistant.const import (
|
|||||||
MAX_LENGTH_STATE_DOMAIN,
|
MAX_LENGTH_STATE_DOMAIN,
|
||||||
MAX_LENGTH_STATE_ENTITY_ID,
|
MAX_LENGTH_STATE_ENTITY_ID,
|
||||||
STATE_UNAVAILABLE,
|
STATE_UNAVAILABLE,
|
||||||
|
STATE_UNKNOWN,
|
||||||
)
|
)
|
||||||
from homeassistant.core import (
|
from homeassistant.core import (
|
||||||
Event,
|
Event,
|
||||||
@ -484,7 +485,7 @@ class EntityRegistry:
|
|||||||
)
|
)
|
||||||
|
|
||||||
@callback
|
@callback
|
||||||
def async_update_entity(
|
def _async_update_entity(
|
||||||
self,
|
self,
|
||||||
entity_id: str,
|
entity_id: str,
|
||||||
*,
|
*,
|
||||||
@ -505,6 +506,8 @@ class EntityRegistry:
|
|||||||
original_name: str | None | UndefinedType = UNDEFINED,
|
original_name: str | None | UndefinedType = UNDEFINED,
|
||||||
supported_features: int | UndefinedType = UNDEFINED,
|
supported_features: int | UndefinedType = UNDEFINED,
|
||||||
unit_of_measurement: str | None | UndefinedType = UNDEFINED,
|
unit_of_measurement: str | None | UndefinedType = UNDEFINED,
|
||||||
|
platform: str | None | UndefinedType = UNDEFINED,
|
||||||
|
options: Mapping[str, Mapping[str, Any]] | UndefinedType = UNDEFINED,
|
||||||
) -> RegistryEntry:
|
) -> RegistryEntry:
|
||||||
"""Private facing update properties method."""
|
"""Private facing update properties method."""
|
||||||
old = self.entities[entity_id]
|
old = self.entities[entity_id]
|
||||||
@ -544,6 +547,8 @@ class EntityRegistry:
|
|||||||
("original_name", original_name),
|
("original_name", original_name),
|
||||||
("supported_features", supported_features),
|
("supported_features", supported_features),
|
||||||
("unit_of_measurement", unit_of_measurement),
|
("unit_of_measurement", unit_of_measurement),
|
||||||
|
("platform", platform),
|
||||||
|
("options", options),
|
||||||
):
|
):
|
||||||
if value is not UNDEFINED and value != getattr(old, attr_name):
|
if value is not UNDEFINED and value != getattr(old, attr_name):
|
||||||
new_values[attr_name] = value
|
new_values[attr_name] = value
|
||||||
@ -595,6 +600,87 @@ class EntityRegistry:
|
|||||||
|
|
||||||
return new
|
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
|
@callback
|
||||||
def async_update_entity_options(
|
def async_update_entity_options(
|
||||||
self, entity_id: str, domain: str, options: dict[str, Any]
|
self, entity_id: str, domain: str, options: dict[str, Any]
|
||||||
@ -602,19 +688,7 @@ class EntityRegistry:
|
|||||||
"""Update entity options."""
|
"""Update entity options."""
|
||||||
old = self.entities[entity_id]
|
old = self.entities[entity_id]
|
||||||
new_options: Mapping[str, Mapping[str, Any]] = {**old.options, domain: options}
|
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:
|
async def async_load(self) -> None:
|
||||||
"""Load the entity registry."""
|
"""Load the entity registry."""
|
||||||
|
@ -1247,3 +1247,74 @@ async def test_entity_category_str_not_allowed(hass):
|
|||||||
reg.async_update_entity(
|
reg.async_update_entity(
|
||||||
entity_id, entity_category=EntityCategory.DIAGNOSTIC.value
|
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,
|
||||||
|
)
|
||||||
|
Loading…
x
Reference in New Issue
Block a user