diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index f10edf1f57d..5529d78e13a 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, Final, Literal, NotRequired, TypedDict +from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypedDict import attr import voluptuous as vol @@ -85,8 +85,6 @@ 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) } @@ -166,17 +164,6 @@ 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.""" @@ -427,17 +414,15 @@ class DeletedRegistryEntry: config_subentry_id: str | None = attr.ib() created_at: datetime = attr.ib() device_class: str | None = attr.ib() - disabled_by: RegistryEntryDisabler | UndefinedType | None = attr.ib() + disabled_by: RegistryEntryDisabler | None = attr.ib() domain: str = attr.ib(init=False, repr=False) - hidden_by: RegistryEntryHider | UndefinedType | None = attr.ib() + hidden_by: RegistryEntryHider | 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 | UndefinedType = attr.ib( - converter=_protect_optional_entity_options - ) + options: ReadOnlyEntityOptionsType = attr.ib(converter=_protect_entity_options) orphaned_timestamp: float | None = attr.ib() _cache: dict[str, Any] = attr.ib(factory=dict, eq=False, init=False) @@ -460,21 +445,15 @@ class DeletedRegistryEntry: "config_subentry_id": self.config_subentry_id, "created_at": self.created_at, "device_class": self.device_class, - "disabled_by": self.disabled_by - if self.disabled_by is not UNDEFINED - else UNDEFINED_STR, + "disabled_by": self.disabled_by, "entity_id": self.entity_id, - "hidden_by": self.hidden_by - if self.hidden_by is not UNDEFINED - else UNDEFINED_STR, + "hidden_by": self.hidden_by, "icon": self.icon, "id": self.id, "labels": list(self.labels), "modified_at": self.modified_at, "name": self.name, - "options": self.options - if self.options is not UNDEFINED - else UNDEFINED_STR, + "options": self.options, "orphaned_timestamp": self.orphaned_timestamp, "platform": self.platform, "unique_id": self.unique_id, @@ -605,12 +584,12 @@ class EntityRegistryStore(storage.Store[dict[str, list[dict[str, Any]]]]): entity["area_id"] = None entity["categories"] = {} entity["device_class"] = None - entity["disabled_by"] = UNDEFINED_STR - entity["hidden_by"] = UNDEFINED_STR + entity["disabled_by"] = None + entity["hidden_by"] = None entity["icon"] = None entity["labels"] = [] entity["name"] = None - entity["options"] = UNDEFINED_STR + entity["options"] = {} if old_major_version > 1: raise NotImplementedError @@ -980,30 +959,25 @@ class EntityRegistry(BaseRegistry): categories = deleted_entity.categories created_at = deleted_entity.created_at device_class = deleted_entity.device_class - 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 + 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 # 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 - if deleted_entity.hidden_by is not UNDEFINED: - hidden_by = deleted_entity.hidden_by + hidden_by = deleted_entity.hidden_by icon = deleted_entity.icon labels = deleted_entity.labels name = deleted_entity.name - if deleted_entity.options is not UNDEFINED: - options = deleted_entity.options - else: - options = get_initial_options() if get_initial_options else None + options = deleted_entity.options else: aliases = set() area_id = None @@ -1556,20 +1530,6 @@ 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] @@ -1587,7 +1547,6 @@ class EntityRegistry(BaseRegistry): entity["platform"], entity["unique_id"], ) - deleted_entities[key] = DeletedRegistryEntry( aliases=set(entity["aliases"]), area_id=entity["area_id"], @@ -1596,21 +1555,23 @@ class EntityRegistry(BaseRegistry): config_subentry_id=entity["config_subentry_id"], created_at=datetime.fromisoformat(entity["created_at"]), device_class=entity["device_class"], - disabled_by=get_optional_enum( - RegistryEntryDisabler, entity["disabled_by"] + disabled_by=( + RegistryEntryDisabler(entity["disabled_by"]) + if entity["disabled_by"] + else None ), entity_id=entity["entity_id"], - hidden_by=get_optional_enum( - RegistryEntryHider, entity["hidden_by"] + hidden_by=( + RegistryEntryHider(entity["hidden_by"]) + if entity["hidden_by"] + else None ), icon=entity["icon"], id=entity["id"], labels=set(entity["labels"]), modified_at=datetime.fromisoformat(entity["modified_at"]), name=entity["name"], - options=entity["options"] - if entity["options"] is not UNDEFINED_STR - else UNDEFINED, + options=entity["options"], 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 da6cdf806d7..acbcb02a5de 100644 --- a/tests/helpers/test_entity_registry.py +++ b/tests/helpers/test_entity_registry.py @@ -20,7 +20,6 @@ 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 ( @@ -963,10 +962,9 @@ 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 migrated data + # Check we store migrated data await flush_store(registry._store) - migrated_data = hass_storage[er.STORAGE_KEY] - assert migrated_data == { + assert hass_storage[er.STORAGE_KEY] == { "version": er.STORAGE_VERSION_MAJOR, "minor_version": er.STORAGE_VERSION_MINOR, "key": er.STORAGE_KEY, @@ -1009,11 +1007,6 @@ 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: @@ -1149,17 +1142,9 @@ 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) - migrated_data = hass_storage[er.STORAGE_KEY] - assert migrated_data == { + assert hass_storage[er.STORAGE_KEY] == { "version": er.STORAGE_VERSION_MAJOR, "minor_version": er.STORAGE_VERSION_MINOR, "key": er.STORAGE_KEY, @@ -1207,15 +1192,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": "UNDEFINED", + "disabled_by": None, "entity_id": "test.deleted_entity", - "hidden_by": "UNDEFINED", + "hidden_by": None, "icon": None, "id": "23456", "labels": [], "modified_at": "1970-01-01T00:00:00+00:00", "name": None, - "options": "UNDEFINED", + "options": {}, "orphaned_timestamp": None, "platform": "super_duper_platform", "unique_id": "very_very_unique", @@ -1224,11 +1209,6 @@ 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 @@ -3170,366 +3150,6 @@ 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",