From 8938c109c25ae072452302fa31a436d4d102eb11 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Mon, 19 May 2025 16:09:23 +0200 Subject: [PATCH] Improve entity registry restore test (#145220) --- tests/helpers/test_entity_registry.py | 142 ++++++++++++++++++++++---- 1 file changed, 122 insertions(+), 20 deletions(-) diff --git a/tests/helpers/test_entity_registry.py b/tests/helpers/test_entity_registry.py index 7df7bb398e8..671c2ddeb29 100644 --- a/tests/helpers/test_entity_registry.py +++ b/tests/helpers/test_entity_registry.py @@ -19,7 +19,7 @@ from homeassistant.const import ( from homeassistant.core import CoreState, HomeAssistant, callback from homeassistant.exceptions import MaxLengthExceeded from homeassistant.helpers import device_registry as dr, entity_registry as er -from homeassistant.util.dt import utc_from_timestamp +from homeassistant.util.dt import utc_from_timestamp, utcnow from tests.common import ( ANY, @@ -2440,10 +2440,11 @@ def test_migrate_entity_to_new_platform_error_handling( async def test_restore_entity( hass: HomeAssistant, + device_registry: dr.DeviceRegistry, entity_registry: er.EntityRegistry, freezer: FrozenDateTimeFactory, ) -> None: - """Make sure entity registry id is stable and entity_id is reused if possible.""" + """Make sure entity registry id is stable.""" update_events = async_capture_events(hass, er.EVENT_ENTITY_REGISTRY_UPDATED) config_entry = MockConfigEntry( domain="light", @@ -2455,11 +2456,44 @@ async def test_restore_entity( title="Mock title", unique_id="test", ), + config_entries.ConfigSubentryData( + data={}, + subentry_id="mock-subentry-id-1-2", + subentry_type="test", + title="Mock title", + unique_id="test", + ), ], ) config_entry.add_to_hass(hass) + device_entry_1 = device_registry.async_get_or_create( + config_entry_id=config_entry.entry_id, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + ) + device_entry_2 = device_registry.async_get_or_create( + config_entry_id=config_entry.entry_id, + connections={(dr.CONNECTION_NETWORK_MAC, "22:34:56:AB:CD:EF")}, + ) entry1 = entity_registry.async_get_or_create( - "light", "hue", "1234", config_entry=config_entry + "light", + "hue", + "1234", + capabilities={"key1": "value1"}, + config_entry=config_entry, + config_subentry_id="mock-subentry-id-1-1", + device_id=device_entry_1.id, + disabled_by=er.RegistryEntryDisabler.DEVICE, + 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="suggested_1", + supported_features=1, + translation_key="translation_key_1", + unit_of_measurement="unit_1", ) entry2 = entity_registry.async_get_or_create( "light", @@ -2469,8 +2503,22 @@ async def test_restore_entity( config_subentry_id="mock-subentry-id-1-1", ) + # Apply user customizations entry1 = entity_registry.async_update_entity( - entry1.entity_id, new_entity_id="light.custom_1" + entry1.entity_id, + aliases={"alias1", "alias2"}, + area_id="12345A", + categories={"scope1": "id", "scope2": "id"}, + device_class="device_class_user", + disabled_by=er.RegistryEntryDisabler.USER, + hidden_by=er.RegistryEntryHider.USER, + icon="icon_user", + labels={"label1", "label2"}, + name="Test Friendly Name", + new_entity_id="light.custom_1", + ) + entry1 = entity_registry.async_update_entity_options( + entry1.entity_id, "options_domain", {"key": "value"} ) entity_registry.async_remove(entry1.entity_id) @@ -2478,17 +2526,61 @@ async def test_restore_entity( assert len(entity_registry.entities) == 0 assert len(entity_registry.deleted_entities) == 2 - # Re-add entities + # Re-add entities, integration has changed entry1_restored = entity_registry.async_get_or_create( - "light", "hue", "1234", config_entry=config_entry + "light", + "hue", + "1234", + capabilities={"key2": "value2"}, + config_entry=config_entry, + config_subentry_id="mock-subentry-id-1-2", + device_id=device_entry_2.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", ) entry2_restored = entity_registry.async_get_or_create("light", "hue", "5678") assert len(entity_registry.entities) == 2 assert len(entity_registry.deleted_entities) == 0 assert entry1 != entry1_restored - # entity_id is not restored - assert attr.evolve(entry1, entity_id="light.hue_1234") == entry1_restored + # entity_id and user customizations are not restored. new integration options are + # respected. + assert entry1_restored == er.RegistryEntry( + entity_id="light.suggested_2", + unique_id="1234", + platform="hue", + capabilities={"key2": "value2"}, + config_entry_id=config_entry.entry_id, + config_subentry_id="mock-subentry-id-1-2", + created_at=utcnow(), + device_class=None, + device_id=device_entry_2.id, + disabled_by=er.RegistryEntryDisabler.INTEGRATION, + entity_category=EntityCategory.CONFIG, + has_entity_name=False, + hidden_by=None, + icon=None, + id=entry1.id, + 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", + supported_features=2, + translation_key="translation_key_2", + unit_of_measurement="unit_2", + ) assert entry2 != entry2_restored # Config entry and subentry are not restored assert ( @@ -2534,23 +2626,33 @@ async def test_restore_entity( # Check the events await hass.async_block_till_done() - assert len(update_events) == 13 - assert update_events[0].data == {"action": "create", "entity_id": "light.hue_1234"} + assert len(update_events) == 14 + assert update_events[0].data == { + "action": "create", + "entity_id": "light.suggested_1", + } assert update_events[1].data == {"action": "create", "entity_id": "light.hue_5678"} assert update_events[2].data["action"] == "update" - assert update_events[3].data == {"action": "remove", "entity_id": "light.custom_1"} - assert update_events[4].data == {"action": "remove", "entity_id": "light.hue_5678"} + assert update_events[3].data["action"] == "update" + assert update_events[4].data == {"action": "remove", "entity_id": "light.custom_1"} + assert update_events[5].data == {"action": "remove", "entity_id": "light.hue_5678"} # Restore entities the 1st time - assert update_events[5].data == {"action": "create", "entity_id": "light.hue_1234"} - assert update_events[6].data == {"action": "create", "entity_id": "light.hue_5678"} - assert update_events[7].data == {"action": "remove", "entity_id": "light.hue_1234"} - assert update_events[8].data == {"action": "remove", "entity_id": "light.hue_5678"} + assert update_events[6].data == { + "action": "create", + "entity_id": "light.suggested_2", + } + assert update_events[7].data == {"action": "create", "entity_id": "light.hue_5678"} + assert update_events[8].data == { + "action": "remove", + "entity_id": "light.suggested_2", + } + assert update_events[9].data == {"action": "remove", "entity_id": "light.hue_5678"} # Restore entities the 2nd time - assert update_events[9].data == {"action": "create", "entity_id": "light.hue_1234"} - assert update_events[10].data == {"action": "create", "entity_id": "light.hue_5678"} - assert update_events[11].data == {"action": "remove", "entity_id": "light.hue_1234"} + assert update_events[10].data == {"action": "create", "entity_id": "light.hue_1234"} + assert update_events[11].data == {"action": "create", "entity_id": "light.hue_5678"} + assert update_events[12].data == {"action": "remove", "entity_id": "light.hue_1234"} # Restore entities the 3rd time - assert update_events[12].data == {"action": "create", "entity_id": "light.hue_1234"} + assert update_events[13].data == {"action": "create", "entity_id": "light.hue_1234"} async def test_async_migrate_entry_delete_self(