From 10fd26df4be168b5373a7ef3e095bf71c9c04900 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Mon, 16 Oct 2023 21:42:02 -1000 Subject: [PATCH] Preserve HomeKit Accessory ID when entity unique id changes (#102123) --- .../components/homekit/aidmanager.py | 27 ++++++++++----- .../components/homekit/type_triggers.py | 4 ++- tests/components/homekit/test_aidmanager.py | 34 +++++++++++++++++-- tests/components/homekit/test_homekit.py | 2 ++ 4 files changed, 55 insertions(+), 12 deletions(-) diff --git a/homeassistant/components/homekit/aidmanager.py b/homeassistant/components/homekit/aidmanager.py index 0deb4500197..43beaaa8dc6 100644 --- a/homeassistant/components/homekit/aidmanager.py +++ b/homeassistant/components/homekit/aidmanager.py @@ -33,9 +33,9 @@ AID_MIN = 2 AID_MAX = 18446744073709551615 -def get_system_unique_id(entity: er.RegistryEntry) -> str: +def get_system_unique_id(entity: er.RegistryEntry, entity_unique_id: str) -> str: """Determine the system wide unique_id for an entity.""" - return f"{entity.platform}.{entity.domain}.{entity.unique_id}" + return f"{entity.platform}.{entity.domain}.{entity_unique_id}" def _generate_aids(unique_id: str | None, entity_id: str) -> Generator[int, None, None]: @@ -72,31 +72,42 @@ class AccessoryAidStorage: self.allocated_aids: set[int] = set() self._entry_id = entry_id self.store: Store | None = None - self._entity_registry: er.EntityRegistry | None = None + self._entity_registry = er.async_get(hass) async def async_initialize(self) -> None: """Load the latest AID data.""" - self._entity_registry = er.async_get(self.hass) aidstore = get_aid_storage_filename_for_entry_id(self._entry_id) self.store = Store(self.hass, AID_MANAGER_STORAGE_VERSION, aidstore) if not (raw_storage := await self.store.async_load()): # There is no data about aid allocations yet return - assert isinstance(raw_storage, dict) self.allocations = raw_storage.get(ALLOCATIONS_KEY, {}) self.allocated_aids = set(self.allocations.values()) def get_or_allocate_aid_for_entity_id(self, entity_id: str) -> int: """Generate a stable aid for an entity id.""" - assert self._entity_registry is not None - if not (entity := self._entity_registry.async_get(entity_id)): + if not (entry := self._entity_registry.async_get(entity_id)): return self.get_or_allocate_aid(None, entity_id) - sys_unique_id = get_system_unique_id(entity) + sys_unique_id = get_system_unique_id(entry, entry.unique_id) + self._migrate_unique_id_aid_assignment_if_needed(sys_unique_id, entry) return self.get_or_allocate_aid(sys_unique_id, entity_id) + def _migrate_unique_id_aid_assignment_if_needed( + self, sys_unique_id: str, entry: er.RegistryEntry + ) -> None: + """Migrate the unique id aid assignment if its changed.""" + if sys_unique_id in self.allocations or not ( + previous_unique_id := entry.previous_unique_id + ): + return + old_sys_unique_id = get_system_unique_id(entry, previous_unique_id) + if aid := self.allocations.pop(old_sys_unique_id, None): + self.allocations[sys_unique_id] = aid + self.async_schedule_save() + def get_or_allocate_aid(self, unique_id: str | None, entity_id: str) -> int: """Allocate (and return) a new aid for an accessory.""" if unique_id and unique_id in self.allocations: diff --git a/homeassistant/components/homekit/type_triggers.py b/homeassistant/components/homekit/type_triggers.py index be8db07d517..8cd01638679 100644 --- a/homeassistant/components/homekit/type_triggers.py +++ b/homeassistant/components/homekit/type_triggers.py @@ -51,7 +51,9 @@ class DeviceTriggerAccessory(HomeAccessory): if (entity_id_or_uuid := trigger.get("entity_id")) and ( entry := ent_reg.async_get(entity_id_or_uuid) ): - unique_id += f"-entity_unique_id:{get_system_unique_id(entry)}" + unique_id += ( + f"-entity_unique_id:{get_system_unique_id(entry, entry.unique_id)}" + ) entity_id = entry.entity_id trigger_name_parts = [] if entity_id and (state := self.hass.states.get(entity_id)): diff --git a/tests/components/homekit/test_aidmanager.py b/tests/components/homekit/test_aidmanager.py index 18e654cb4ed..447cdc99a57 100644 --- a/tests/components/homekit/test_aidmanager.py +++ b/tests/components/homekit/test_aidmanager.py @@ -66,9 +66,9 @@ async def test_aid_generation( == 1751603975 ) - aid_storage.delete_aid(get_system_unique_id(light_ent)) - aid_storage.delete_aid(get_system_unique_id(light_ent2)) - aid_storage.delete_aid(get_system_unique_id(remote_ent)) + aid_storage.delete_aid(get_system_unique_id(light_ent, light_ent.unique_id)) + aid_storage.delete_aid(get_system_unique_id(light_ent2, light_ent2.unique_id)) + aid_storage.delete_aid(get_system_unique_id(remote_ent, remote_ent.unique_id)) aid_storage.delete_aid("non-existent-one") for _ in range(0, 2): @@ -618,3 +618,31 @@ async def test_aid_generation_no_unique_ids_handles_collision( aid_storage_path = hass.config.path(STORAGE_DIR, aidstore) if await hass.async_add_executor_job(os.path.exists, aid_storage_path): await hass.async_add_executor_job(os.unlink, aid_storage_path) + + +async def test_handle_unique_id_change( + hass: HomeAssistant, +) -> None: + """Test handling unique id changes.""" + entity_registry = er.async_get(hass) + light = entity_registry.async_get_or_create("light", "demo", "old_unique") + config_entry = MockConfigEntry(domain="test", data={}) + config_entry.add_to_hass(hass) + with patch( + "homeassistant.components.homekit.aidmanager.AccessoryAidStorage.async_schedule_save" + ): + aid_storage = AccessoryAidStorage(hass, config_entry) + await aid_storage.async_initialize() + + original_aid = aid_storage.get_or_allocate_aid_for_entity_id(light.entity_id) + assert aid_storage.allocations == {"demo.light.old_unique": 4202023227} + + entity_registry.async_update_entity(light.entity_id, new_unique_id="new_unique") + await hass.async_block_till_done() + + aid = aid_storage.get_or_allocate_aid_for_entity_id(light.entity_id) + assert aid == original_aid + + # Verify that the old unique id is removed from the allocations + # and that the new unique id assumes the old aid + assert aid_storage.allocations == {"demo.light.new_unique": 4202023227} diff --git a/tests/components/homekit/test_homekit.py b/tests/components/homekit/test_homekit.py index 5c517ac9cb9..158efa477d4 100644 --- a/tests/components/homekit/test_homekit.py +++ b/tests/components/homekit/test_homekit.py @@ -841,6 +841,7 @@ async def test_homekit_start_with_a_device( await async_init_entry(hass, entry) homekit = _mock_homekit(hass, entry, HOMEKIT_MODE_BRIDGE, None, devices=[device_id]) homekit.driver = hk_driver + homekit.aid_storage = MagicMock() with patch(f"{PATH_HOMEKIT}.get_accessory", side_effect=Exception), patch( f"{PATH_HOMEKIT}.async_show_setup_message" @@ -868,6 +869,7 @@ async def test_homekit_stop(hass: HomeAssistant) -> None: homekit.driver.async_stop = AsyncMock() homekit.bridge = Mock() homekit.bridge.accessories = {} + homekit.aid_storage = MagicMock() assert homekit.status == STATUS_READY await homekit.async_stop()