home-assistant/core
1
0
Fork 1
mirror of https://github.com/home-assistant/core.git synced 2025-08-30 18:01:31 +02:00
Code Issues Packages Projects Releases Wiki Activity

Improve migration to entity registry version 1.18 (#151308)

Browse Source 
Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
This commit is contained in:
Erik Montnemery
2025-08-28 18:52:51 +02:00
committed by GitHub
parent ffcd5167b5
commit 4130f3db2f
2 changed files with 454 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
homeassistant/helpers/entity_registry.py
View File

@@ -16,7 +16,7 @@ from datetime import datetime, timedelta
from enum import StrEnum
import logging
import time
from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypedDict
from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, TypedDict
import attr
import voluptuous as vol
@@ -85,6 +85,8 @@ STORAGE_KEY = "core.entity_registry"
CLEANUP_INTERVAL = 3600 * 24
ORPHANED_ENTITY_KEEP_SECONDS = 3600 * 24 * 30
UNDEFINED_STR: Final = "UNDEFINED"
ENTITY_CATEGORY_VALUE_TO_INDEX: dict[EntityCategory | None, int] = {
val: idx for idx, val in enumerate(EntityCategory)
}
@@ -164,6 +166,17 @@ def _protect_entity_options(
return ReadOnlyDict({key: ReadOnlyDict(val) for key, val in data.items()})
def _protect_optional_entity_options(
data: EntityOptionsType | UndefinedType | None,
) -> ReadOnlyEntityOptionsType | UndefinedType:
"""Protect entity options from being modified."""
if data is UNDEFINED:
return UNDEFINED
if data is None:
return ReadOnlyDict({})
return ReadOnlyDict({key: ReadOnlyDict(val) for key, val in data.items()})
@attr.s(frozen=True, kw_only=True, slots=True)
class RegistryEntry:
"""Entity Registry Entry."""
@@ -414,15 +427,17 @@ class DeletedRegistryEntry:
config_subentry_id: str | None = attr.ib()
created_at: datetime = attr.ib()
device_class: str | None = attr.ib()
disabled_by: RegistryEntryDisabler | None = attr.ib()
disabled_by: RegistryEntryDisabler | UndefinedType | None = attr.ib()
domain: str = attr.ib(init=False, repr=False)
hidden_by: RegistryEntryHider | None = attr.ib()
hidden_by: RegistryEntryHider | UndefinedType | None = attr.ib()
icon: str | None = attr.ib()
id: str = attr.ib()
labels: set[str] = attr.ib()
modified_at: datetime = attr.ib()
name: str | None = attr.ib()
options: ReadOnlyEntityOptionsType = attr.ib(converter=_protect_entity_options)
options: ReadOnlyEntityOptionsType | UndefinedType = attr.ib(
converter=_protect_optional_entity_options
)
orphaned_timestamp: float | None = attr.ib()
_cache: dict[str, Any] = attr.ib(factory=dict, eq=False, init=False)
@@ -445,15 +460,21 @@ class DeletedRegistryEntry:
"config_subentry_id": self.config_subentry_id,
"created_at": self.created_at,
"device_class": self.device_class,
"disabled_by": self.disabled_by,
"disabled_by": self.disabled_by
if self.disabled_by is not UNDEFINED
else UNDEFINED_STR,
"entity_id": self.entity_id,
"hidden_by": self.hidden_by,
"hidden_by": self.hidden_by
if self.hidden_by is not UNDEFINED
else UNDEFINED_STR,
"icon": self.icon,
"id": self.id,
"labels": list(self.labels),
"modified_at": self.modified_at,
"name": self.name,
"options": self.options,
"options": self.options
if self.options is not UNDEFINED
else UNDEFINED_STR,
"orphaned_timestamp": self.orphaned_timestamp,
"platform": self.platform,
"unique_id": self.unique_id,
@@ -584,12 +605,12 @@ class EntityRegistryStore(storage.Store[dict[str, list[dict[str, Any]]]]):
entity["area_id"] = None
entity["categories"] = {}
entity["device_class"] = None
entity["disabled_by"] = None
entity["hidden_by"] = None
entity["disabled_by"] = UNDEFINED_STR
entity["hidden_by"] = UNDEFINED_STR
entity["icon"] = None
entity["labels"] = []
entity["name"] = None
entity["options"] = {}
entity["options"] = UNDEFINED_STR
if old_major_version > 1:
raise NotImplementedError
@@ -958,25 +979,30 @@ class EntityRegistry(BaseRegistry):
categories = deleted_entity.categories
created_at = deleted_entity.created_at
device_class = deleted_entity.device_class
disabled_by = deleted_entity.disabled_by
# Adjust disabled_by based on config entry state
if config_entry and config_entry is not UNDEFINED:
if config_entry.disabled_by:
if disabled_by is None:
disabled_by = RegistryEntryDisabler.CONFIG_ENTRY
if deleted_entity.disabled_by is not UNDEFINED:
disabled_by = deleted_entity.disabled_by
# Adjust disabled_by based on config entry state
if config_entry and config_entry is not UNDEFINED:
if config_entry.disabled_by:
if disabled_by is None:
disabled_by = RegistryEntryDisabler.CONFIG_ENTRY
elif disabled_by == RegistryEntryDisabler.CONFIG_ENTRY:
disabled_by = None
elif disabled_by == RegistryEntryDisabler.CONFIG_ENTRY:
disabled_by = None
elif disabled_by == RegistryEntryDisabler.CONFIG_ENTRY:
disabled_by = None
# Restore entity_id if it's available
if self._entity_id_available(deleted_entity.entity_id):
entity_id = deleted_entity.entity_id
entity_registry_id = deleted_entity.id
hidden_by = deleted_entity.hidden_by
if deleted_entity.hidden_by is not UNDEFINED:
hidden_by = deleted_entity.hidden_by
icon = deleted_entity.icon
labels = deleted_entity.labels
name = deleted_entity.name
options = deleted_entity.options
if deleted_entity.options is not UNDEFINED:
options = deleted_entity.options
else:
options = get_initial_options() if get_initial_options else None
else:
aliases = set()
area_id = None
@@ -1529,6 +1555,20 @@ class EntityRegistry(BaseRegistry):
previous_unique_id=entity["previous_unique_id"],
unit_of_measurement=entity["unit_of_measurement"],
)
def get_optional_enum[_EnumT: StrEnum](
cls: type[_EnumT], value: str | None
) -> _EnumT | UndefinedType | None:
"""Convert string to the passed enum, UNDEFINED or None."""
if value is None:
return None
if value == UNDEFINED_STR:
return UNDEFINED
try:
return cls(value)
except ValueError:
return None
for entity in data["deleted_entities"]:
try:
domain = split_entity_id(entity["entity_id"])[0]
@@ -1546,6 +1586,7 @@ class EntityRegistry(BaseRegistry):
entity["platform"],
entity["unique_id"],
)
deleted_entities[key] = DeletedRegistryEntry(
aliases=set(entity["aliases"]),
area_id=entity["area_id"],
@@ -1554,23 +1595,21 @@ class EntityRegistry(BaseRegistry):
config_subentry_id=entity["config_subentry_id"],
created_at=datetime.fromisoformat(entity["created_at"]),
device_class=entity["device_class"],
disabled_by=(
RegistryEntryDisabler(entity["disabled_by"])
if entity["disabled_by"]
else None
disabled_by=get_optional_enum(
RegistryEntryDisabler, entity["disabled_by"]
),
entity_id=entity["entity_id"],
hidden_by=(
RegistryEntryHider(entity["hidden_by"])
if entity["hidden_by"]
else None
hidden_by=get_optional_enum(
RegistryEntryHider, entity["hidden_by"]
),
icon=entity["icon"],
id=entity["id"],
labels=set(entity["labels"]),
modified_at=datetime.fromisoformat(entity["modified_at"]),
name=entity["name"],
options=entity["options"],
options=entity["options"]
if entity["options"] is not UNDEFINED_STR
else UNDEFINED,
orphaned_timestamp=entity["orphaned_timestamp"],
platform=entity["platform"],
unique_id=entity["unique_id"],

392
tests/helpers/test_entity_registry.py
View File

@@ -20,6 +20,7 @@ from homeassistant.core import CoreState, Event, HomeAssistant, callback
from homeassistant.exceptions import MaxLengthExceeded
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.event import async_track_entity_registry_updated_event
from homeassistant.helpers.typing import UNDEFINED
from homeassistant.util.dt import utc_from_timestamp, utcnow
from tests.common import (
@@ -962,9 +963,10 @@ async def test_migration_1_1(hass: HomeAssistant, hass_storage: dict[str, Any])
assert entry.device_class is None
assert entry.original_device_class == "best_class"
# Check we store migrated data
# Check migrated data
await flush_store(registry._store)
assert hass_storage[er.STORAGE_KEY] == {
migrated_data = hass_storage[er.STORAGE_KEY]
assert migrated_data == {
"version": er.STORAGE_VERSION_MAJOR,
"minor_version": er.STORAGE_VERSION_MINOR,
"key": er.STORAGE_KEY,
@@ -1007,6 +1009,11 @@ async def test_migration_1_1(hass: HomeAssistant, hass_storage: dict[str, Any])
},
}
# Serialize the migrated data again
registry.async_schedule_save()
await flush_store(registry._store)
assert hass_storage[er.STORAGE_KEY] == migrated_data
@pytest.mark.parametrize("load_registries", [False])
async def test_migration_1_7(hass: HomeAssistant, hass_storage: dict[str, Any]) -> None:
@@ -1142,9 +1149,17 @@ async def test_migration_1_11(
assert entry.device_class is None
assert entry.original_device_class == "best_class"
deleted_entry = registry.deleted_entities[
("test", "super_duper_platform", "very_very_unique")
]
assert deleted_entry.disabled_by is UNDEFINED
assert deleted_entry.hidden_by is UNDEFINED
assert deleted_entry.options is UNDEFINED
# Check migrated data
await flush_store(registry._store)
assert hass_storage[er.STORAGE_KEY] == {
migrated_data = hass_storage[er.STORAGE_KEY]
assert migrated_data == {
"version": er.STORAGE_VERSION_MAJOR,
"minor_version": er.STORAGE_VERSION_MINOR,
"key": er.STORAGE_KEY,
@@ -1192,15 +1207,15 @@ async def test_migration_1_11(
"config_subentry_id": None,
"created_at": "1970-01-01T00:00:00+00:00",
"device_class": None,
"disabled_by": None,
"disabled_by": "UNDEFINED",
"entity_id": "test.deleted_entity",
"hidden_by": None,
"hidden_by": "UNDEFINED",
"icon": None,
"id": "23456",
"labels": [],
"modified_at": "1970-01-01T00:00:00+00:00",
"name": None,
"options": {},
"options": "UNDEFINED",
"orphaned_timestamp": None,
"platform": "super_duper_platform",
"unique_id": "very_very_unique",
@@ -1209,6 +1224,11 @@ async def test_migration_1_11(
},
}
# Serialize the migrated data again
registry.async_schedule_save()
await flush_store(registry._store)
assert hass_storage[er.STORAGE_KEY] == migrated_data
async def test_update_entity_unique_id(
hass: HomeAssistant, entity_registry: er.EntityRegistry
@@ -3150,6 +3170,366 @@ async def test_restore_entity(
assert update_events[16].data == {"action": "create", "entity_id": "light.hue_1234"}
@pytest.mark.parametrize(
("entity_disabled_by"),
[
None,
er.RegistryEntryDisabler.CONFIG_ENTRY,
er.RegistryEntryDisabler.DEVICE,
er.RegistryEntryDisabler.HASS,
er.RegistryEntryDisabler.INTEGRATION,
er.RegistryEntryDisabler.USER,
],
)
@pytest.mark.usefixtures("freezer")
async def test_restore_migrated_entity_disabled_by(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
entity_disabled_by: er.RegistryEntryDisabler | None,
) -> None:
"""Check how the disabled_by flag is treated when restoring an entity."""
update_events = async_capture_events(hass, er.EVENT_ENTITY_REGISTRY_UPDATED)
config_entry = MockConfigEntry(domain="light")
config_entry.add_to_hass(hass)
device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
)
entry = entity_registry.async_get_or_create(
"light",
"hue",
"1234",
capabilities={"key1": "value1"},
config_entry=config_entry,
config_subentry_id=None,
device_id=device_entry.id,
disabled_by=None,
entity_category=EntityCategory.DIAGNOSTIC,
get_initial_options=lambda: {"test_domain": {"key1": "value1"}},
has_entity_name=True,
hidden_by=er.RegistryEntryHider.INTEGRATION,
original_device_class="device_class_1",
original_icon="original_icon_1",
original_name="original_name_1",
suggested_object_id="hue_5678",
supported_features=1,
translation_key="translation_key_1",
unit_of_measurement="unit_1",
)
entity_registry.async_remove(entry.entity_id)
assert len(entity_registry.entities) == 0
assert len(entity_registry.deleted_entities) == 1
deleted_entry = entity_registry.deleted_entities[("light", "hue", "1234")]
entity_registry.deleted_entities[("light", "hue", "1234")] = attr.evolve(
deleted_entry, disabled_by=UNDEFINED
)
# Re-add entity, integration has changed
entry_restored = entity_registry.async_get_or_create(
"light",
"hue",
"1234",
capabilities={"key2": "value2"},
config_entry=config_entry,
config_subentry_id=None,
device_id=device_entry.id,
disabled_by=entity_disabled_by,
entity_category=EntityCategory.CONFIG,
get_initial_options=lambda: {"test_domain": {"key2": "value2"}},
has_entity_name=False,
hidden_by=None,
original_device_class="device_class_2",
original_icon="original_icon_2",
original_name="original_name_2",
suggested_object_id="suggested_2",
supported_features=2,
translation_key="translation_key_2",
unit_of_measurement="unit_2",
)
assert len(entity_registry.entities) == 1
assert len(entity_registry.deleted_entities) == 0
assert entry != entry_restored
# entity_id and user customizations are restored. new integration options are
# respected.
assert entry_restored == er.RegistryEntry(
entity_id="light.hue_5678",
unique_id="1234",
platform="hue",
aliases=set(),
area_id=None,
categories={},
capabilities={"key2": "value2"},
config_entry_id=config_entry.entry_id,
config_subentry_id=None,
created_at=utcnow(),
device_class=None,
device_id=device_entry.id,
disabled_by=entity_disabled_by,
entity_category=EntityCategory.CONFIG,
has_entity_name=False,
hidden_by=er.RegistryEntryHider.INTEGRATION,
icon=None,
id=entry.id,
labels=set(),
modified_at=utcnow(),
name=None,
options={"test_domain": {"key1": "value1"}},
original_device_class="device_class_2",
original_icon="original_icon_2",
original_name="original_name_2",
suggested_object_id="suggested_2",
supported_features=2,
translation_key="translation_key_2",
unit_of_measurement="unit_2",
)
# Check the events
await hass.async_block_till_done()
assert len(update_events) == 3
assert update_events[0].data == {"action": "create", "entity_id": "light.hue_5678"}
assert update_events[1].data == {"action": "remove", "entity_id": "light.hue_5678"}
assert update_events[2].data == {"action": "create", "entity_id": "light.hue_5678"}
@pytest.mark.parametrize(
("entity_hidden_by"),
[
None,
er.RegistryEntryHider.INTEGRATION,
er.RegistryEntryHider.USER,
],
)
@pytest.mark.usefixtures("freezer")
async def test_restore_migrated_entity_hidden_by(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
entity_hidden_by: er.RegistryEntryHider | None,
) -> None:
"""Check how the hidden_by flag is treated when restoring an entity."""
update_events = async_capture_events(hass, er.EVENT_ENTITY_REGISTRY_UPDATED)
config_entry = MockConfigEntry(domain="light")
config_entry.add_to_hass(hass)
device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
)
entry = entity_registry.async_get_or_create(
"light",
"hue",
"1234",
capabilities={"key1": "value1"},
config_entry=config_entry,
config_subentry_id=None,
device_id=device_entry.id,
disabled_by=None,
entity_category=EntityCategory.DIAGNOSTIC,
get_initial_options=lambda: {"test_domain": {"key1": "value1"}},
has_entity_name=True,
hidden_by=er.RegistryEntryHider.INTEGRATION,
original_device_class="device_class_1",
original_icon="original_icon_1",
original_name="original_name_1",
suggested_object_id="hue_5678",
supported_features=1,
translation_key="translation_key_1",
unit_of_measurement="unit_1",
)
entity_registry.async_remove(entry.entity_id)
assert len(entity_registry.entities) == 0
assert len(entity_registry.deleted_entities) == 1
deleted_entry = entity_registry.deleted_entities[("light", "hue", "1234")]
entity_registry.deleted_entities[("light", "hue", "1234")] = attr.evolve(
deleted_entry, hidden_by=UNDEFINED
)
# Re-add entity, integration has changed
entry_restored = entity_registry.async_get_or_create(
"light",
"hue",
"1234",
capabilities={"key2": "value2"},
config_entry=config_entry,
config_subentry_id=None,
device_id=device_entry.id,
disabled_by=None,
entity_category=EntityCategory.CONFIG,
get_initial_options=lambda: {"test_domain": {"key2": "value2"}},
has_entity_name=False,
hidden_by=entity_hidden_by,
original_device_class="device_class_2",
original_icon="original_icon_2",
original_name="original_name_2",
suggested_object_id="suggested_2",
supported_features=2,
translation_key="translation_key_2",
unit_of_measurement="unit_2",
)
assert len(entity_registry.entities) == 1
assert len(entity_registry.deleted_entities) == 0
assert entry != entry_restored
# entity_id and user customizations are restored. new integration options are
# respected.
assert entry_restored == er.RegistryEntry(
entity_id="light.hue_5678",
unique_id="1234",
platform="hue",
aliases=set(),
area_id=None,
categories={},
capabilities={"key2": "value2"},
config_entry_id=config_entry.entry_id,
config_subentry_id=None,
created_at=utcnow(),
device_class=None,
device_id=device_entry.id,
disabled_by=None,
entity_category=EntityCategory.CONFIG,
has_entity_name=False,
hidden_by=entity_hidden_by,
icon=None,
id=entry.id,
labels=set(),
modified_at=utcnow(),
name=None,
options={"test_domain": {"key1": "value1"}},
original_device_class="device_class_2",
original_icon="original_icon_2",
original_name="original_name_2",
suggested_object_id="suggested_2",
supported_features=2,
translation_key="translation_key_2",
unit_of_measurement="unit_2",
)
# Check the events
await hass.async_block_till_done()
assert len(update_events) == 3
assert update_events[0].data == {"action": "create", "entity_id": "light.hue_5678"}
assert update_events[1].data == {"action": "remove", "entity_id": "light.hue_5678"}
assert update_events[2].data == {"action": "create", "entity_id": "light.hue_5678"}
@pytest.mark.usefixtures("freezer")
async def test_restore_migrated_entity_initial_options(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Check how the initial options is treated when restoring an entity."""
update_events = async_capture_events(hass, er.EVENT_ENTITY_REGISTRY_UPDATED)
config_entry = MockConfigEntry(domain="light")
config_entry.add_to_hass(hass)
device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
)
entry = entity_registry.async_get_or_create(
"light",
"hue",
"1234",
capabilities={"key1": "value1"},
config_entry=config_entry,
config_subentry_id=None,
device_id=device_entry.id,
disabled_by=None,
entity_category=EntityCategory.DIAGNOSTIC,
get_initial_options=lambda: {"test_domain": {"key1": "value1"}},
has_entity_name=True,
hidden_by=er.RegistryEntryHider.INTEGRATION,
original_device_class="device_class_1",
original_icon="original_icon_1",
original_name="original_name_1",
suggested_object_id="hue_5678",
supported_features=1,
translation_key="translation_key_1",
unit_of_measurement="unit_1",
)
entity_registry.async_remove(entry.entity_id)
assert len(entity_registry.entities) == 0
assert len(entity_registry.deleted_entities) == 1
deleted_entry = entity_registry.deleted_entities[("light", "hue", "1234")]
entity_registry.deleted_entities[("light", "hue", "1234")] = attr.evolve(
deleted_entry, options=UNDEFINED
)
# Re-add entity, integration has changed
entry_restored = entity_registry.async_get_or_create(
"light",
"hue",
"1234",
capabilities={"key2": "value2"},
config_entry=config_entry,
config_subentry_id=None,
device_id=device_entry.id,
disabled_by=None,
entity_category=EntityCategory.CONFIG,
get_initial_options=lambda: {"test_domain": {"key2": "value2"}},
has_entity_name=False,
hidden_by=None,
original_device_class="device_class_2",
original_icon="original_icon_2",
original_name="original_name_2",
suggested_object_id="suggested_2",
supported_features=2,
translation_key="translation_key_2",
unit_of_measurement="unit_2",
)
assert len(entity_registry.entities) == 1
assert len(entity_registry.deleted_entities) == 0
assert entry != entry_restored
# entity_id and user customizations are restored. new integration options are
# respected.
assert entry_restored == er.RegistryEntry(
entity_id="light.hue_5678",
unique_id="1234",
platform="hue",
aliases=set(),
area_id=None,
categories={},
capabilities={"key2": "value2"},
config_entry_id=config_entry.entry_id,
config_subentry_id=None,
created_at=utcnow(),
device_class=None,
device_id=device_entry.id,
disabled_by=None,
entity_category=EntityCategory.CONFIG,
has_entity_name=False,
hidden_by=er.RegistryEntryHider.INTEGRATION,
icon=None,
id=entry.id,
labels=set(),
modified_at=utcnow(),
name=None,
options={"test_domain": {"key2": "value2"}},
original_device_class="device_class_2",
original_icon="original_icon_2",
original_name="original_name_2",
suggested_object_id="suggested_2",
supported_features=2,
translation_key="translation_key_2",
unit_of_measurement="unit_2",
)
# Check the events
await hass.async_block_till_done()
assert len(update_events) == 3
assert update_events[0].data == {"action": "create", "entity_id": "light.hue_5678"}
assert update_events[1].data == {"action": "remove", "entity_id": "light.hue_5678"}
assert update_events[2].data == {"action": "create", "entity_id": "light.hue_5678"}
@pytest.mark.parametrize(
(
"config_entry_disabled_by",