Preserve HomeKit Accessory ID when entity unique id changes (#102123)

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:

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)):

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}

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()