github/ha-core
1
Fork 0
mirror of https://github.com/home-assistant/core synced 2024-09-18 19:55:20 +02:00
Code Releases Activity

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

Browse Source 
* 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:
Raman Gupta 2022-04-16 16:18:52 -04:00 committed by GitHub
parent 42e0bc849c
commit 3bcd921a28
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 172 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
homeassistant/components/tomorrowio/__init__.py
View File

@ -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)
ent_reg.async_remove(entity_entry.entity_id)
old_platform = entity_entry.platform
# 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(
"Re-creating %s for the new config entry", entity_entry.entity_id
)
# We will precreate the entity so that any customizations can be preserved
new_entity_entry = ent_reg.async_get_or_create(
entity_entry.domain,
ent_reg.async_update_entity_platform(
entity_entry.entity_id,
DOMAIN,
new_unique_id=new_unique_id,
new_config_entry_id=entry.entry_id,
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(

102
homeassistant/helpers/entity_registry.py
View File

@ -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)
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
return self._async_update_entity(entity_id, options=new_options)
async def async_load(self) -> None:
"""Load the entity registry."""

71
tests/helpers/test_entity_registry.py
View File

@ -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,
)