From 05c8e8b4fd73584e5346c14fa65d28567b50ffc0 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Mon, 25 Aug 2025 22:47:55 +0200 Subject: [PATCH] Adjust entity disabled_by flag when restoring a deleted entity (#151150) --- homeassistant/helpers/entity_registry.py | 9 + tests/helpers/test_entity_registry.py | 310 +++++++++++++++++++++++ 2 files changed, 319 insertions(+) diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index 2125c0f4512..3b506f9a2cd 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -959,6 +959,15 @@ class EntityRegistry(BaseRegistry): 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 + 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 diff --git a/tests/helpers/test_entity_registry.py b/tests/helpers/test_entity_registry.py index 89822b80039..151bd54dc2a 100644 --- a/tests/helpers/test_entity_registry.py +++ b/tests/helpers/test_entity_registry.py @@ -2990,6 +2990,316 @@ async def test_restore_entity( assert update_events[16].data == {"action": "create", "entity_id": "light.hue_1234"} +@pytest.mark.parametrize( + ( + "config_entry_disabled_by", + "entity_disabled_by_initial", + "entity_disabled_by_restored", + ), + [ + ( + None, + None, + None, + ), + # Config entry not disabled, entity was disabled by config entry. + # Entity not disabled when restored. + ( + None, + er.RegistryEntryDisabler.CONFIG_ENTRY, + None, + ), + ( + None, + er.RegistryEntryDisabler.DEVICE, + er.RegistryEntryDisabler.DEVICE, + ), + ( + None, + er.RegistryEntryDisabler.HASS, + er.RegistryEntryDisabler.HASS, + ), + ( + None, + er.RegistryEntryDisabler.INTEGRATION, + er.RegistryEntryDisabler.INTEGRATION, + ), + ( + None, + er.RegistryEntryDisabler.USER, + er.RegistryEntryDisabler.USER, + ), + # Config entry disabled, entity not disabled. + # Entity disabled by config entry when restored. + ( + config_entries.ConfigEntryDisabler.USER, + None, + er.RegistryEntryDisabler.CONFIG_ENTRY, + ), + ( + config_entries.ConfigEntryDisabler.USER, + er.RegistryEntryDisabler.CONFIG_ENTRY, + er.RegistryEntryDisabler.CONFIG_ENTRY, + ), + ( + config_entries.ConfigEntryDisabler.USER, + er.RegistryEntryDisabler.DEVICE, + er.RegistryEntryDisabler.DEVICE, + ), + ( + config_entries.ConfigEntryDisabler.USER, + er.RegistryEntryDisabler.HASS, + er.RegistryEntryDisabler.HASS, + ), + ( + config_entries.ConfigEntryDisabler.USER, + er.RegistryEntryDisabler.INTEGRATION, + er.RegistryEntryDisabler.INTEGRATION, + ), + ( + config_entries.ConfigEntryDisabler.USER, + er.RegistryEntryDisabler.USER, + er.RegistryEntryDisabler.USER, + ), + ], +) +@pytest.mark.usefixtures("freezer") +async def test_restore_entity_disabled_by( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + config_entry_disabled_by: config_entries.ConfigEntryDisabler | None, + entity_disabled_by_initial: er.RegistryEntryDisabler | None, + entity_disabled_by_restored: 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", disabled_by=config_entry_disabled_by) + 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=entity_disabled_by_initial, + 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 + + # 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=er.RegistryEntryDisabler.INTEGRATION, + 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_restored, + 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_disabled_by_initial", "entity_disabled_by_restored"), + [ + (None, None), + # Config entry not disabled, entity was disabled by config entry. + # Entity not disabled when restored. + (er.RegistryEntryDisabler.CONFIG_ENTRY, None), + (er.RegistryEntryDisabler.DEVICE, er.RegistryEntryDisabler.DEVICE), + (er.RegistryEntryDisabler.HASS, er.RegistryEntryDisabler.HASS), + (er.RegistryEntryDisabler.INTEGRATION, er.RegistryEntryDisabler.INTEGRATION), + (er.RegistryEntryDisabler.USER, er.RegistryEntryDisabler.USER), + ], +) +@pytest.mark.usefixtures("freezer") +async def test_restore_entity_disabled_by_2( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + entity_disabled_by_initial: er.RegistryEntryDisabler | None, + entity_disabled_by_restored: er.RegistryEntryDisabler | None, +) -> None: + """Check how the disabled_by flag is treated when restoring an entity. + + In this test, the entity is restored without a config entry. + """ + 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=entity_disabled_by_initial, + 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 + + # Re-add entity, integration has changed + entry_restored = entity_registry.async_get_or_create( + "light", + "hue", + "1234", + capabilities={"key2": "value2"}, + config_entry=None, + config_subentry_id=None, + device_id=device_entry.id, + disabled_by=er.RegistryEntryDisabler.INTEGRATION, + 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=None, + config_subentry_id=None, + created_at=utcnow(), + device_class=None, + device_id=device_entry.id, + disabled_by=entity_disabled_by_restored, + 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"} + + async def test_async_migrate_entry_delete_self( hass: HomeAssistant, entity_registry: er.EntityRegistry ) -> None: