From 4130f3db2f68b7ba854868b85b45726955349864 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Thu, 28 Aug 2025 18:52:51 +0200 Subject: [PATCH] Improve migration to entity registry version 1.18 (#151308) Co-authored-by: Martin Hjelmare --- homeassistant/helpers/entity_registry.py | 97 ++++-- tests/helpers/test_entity_registry.py | 392 ++++++++++++++++++++++- 2 files changed, 454 insertions(+), 35 deletions(-) diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index 571f914e9d3..95aa153ff00 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -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"], diff --git a/tests/helpers/test_entity_registry.py b/tests/helpers/test_entity_registry.py index acbcb02a5de..da6cdf806d7 100644 --- a/tests/helpers/test_entity_registry.py +++ b/tests/helpers/test_entity_registry.py @@ -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",