Cleanup device registry for tedee when a lock is removed (#106994)

* remove removed locks

* move duplicated code to function

* remove entities by removing device

* add new locks automatically

* add locks from coordinator

* smaller pr

* remove snapshot

* move lock removal to coordinator

* change comment

* Update tests/components/tedee/test_init.py

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

* Update tests/components/tedee/test_init.py

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

* test lock unavailable

* move logic to function

* resolve merge conflicts

* no need to call keys()

* no need to call keys()

* check for change first

* readability

* Update tests/components/tedee/test_lock.py

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

* Update tests/components/tedee/test_lock.py

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

---------

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

homeassistant/components/tedee/coordinator.py

@@ -18,6 +18,7 @@ from homeassistant.const import CONF_HOST
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
+import homeassistant.helpers.device_registry as dr
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import CONF_LOCAL_ACCESS_TOKEN, DOMAIN
@@ -50,7 +51,7 @@ class TedeeApiCoordinator(DataUpdateCoordinator[dict[int, TedeeLock]]):
         )
 
         self._next_get_locks = time.time()
-        self._current_locks: set[int] = set()
+        self._locks_last_update: set[int] = set()
         self.new_lock_callbacks: list[Callable[[int], None]] = []
 
     @property
@@ -84,15 +85,7 @@ class TedeeApiCoordinator(DataUpdateCoordinator[dict[int, TedeeLock]]):
             ", ".join(map(str, self.tedee_client.locks_dict.keys())),
         )
 
-        if not self._current_locks:
+        self._async_add_remove_locks()
-            self._current_locks = set(self.tedee_client.locks_dict.keys())
-
-        if new_locks := set(self.tedee_client.locks_dict.keys()) - self._current_locks:
-            _LOGGER.debug("New locks found: %s", ", ".join(map(str, new_locks)))
-            for lock_id in new_locks:
-                for callback in self.new_lock_callbacks:
-                    callback(lock_id)
-
         return self.tedee_client.locks_dict
 
     async def _async_update(self, update_fn: Callable[[], Awaitable[None]]) -> None:
@@ -109,3 +102,32 @@ class TedeeApiCoordinator(DataUpdateCoordinator[dict[int, TedeeLock]]):
             raise UpdateFailed("Error while updating data: %s" % str(ex)) from ex
         except (TedeeClientException, TimeoutError) as ex:
             raise UpdateFailed("Querying API failed. Error: %s" % str(ex)) from ex
+
+    def _async_add_remove_locks(self) -> None:
+        """Add new locks, remove non-existing locks."""
+        if not self._locks_last_update:
+            self._locks_last_update = set(self.tedee_client.locks_dict)
+
+        if (
+            current_locks := set(self.tedee_client.locks_dict)
+        ) == self._locks_last_update:
+            return
+
+        # remove old locks
+        if removed_locks := self._locks_last_update - current_locks:
+            _LOGGER.debug("Removed locks: %s", ", ".join(map(str, removed_locks)))
+            device_registry = dr.async_get(self.hass)
+            for lock_id in removed_locks:
+                if device := device_registry.async_get_device(
+                    identifiers={(DOMAIN, str(lock_id))}
+                ):
+                    device_registry.async_remove_device(device.id)
+
+        # add new locks
+        if new_locks := current_locks - self._locks_last_update:
+            _LOGGER.debug("New locks found: %s", ", ".join(map(str, new_locks)))
+            for lock_id in new_locks:
+                for callback in self.new_lock_callbacks:
+                    callback(lock_id)
+
+        self._locks_last_update = current_locks

homeassistant/components/tedee/entity.py

@@ -38,9 +38,14 @@ class TedeeEntity(CoordinatorEntity[TedeeApiCoordinator]):
     @callback
     def _handle_coordinator_update(self) -> None:
         """Handle updated data from the coordinator."""
-        self._lock = self.coordinator.data[self._lock.lock_id]
+        self._lock = self.coordinator.data.get(self._lock.lock_id, self._lock)
         super()._handle_coordinator_update()
 
+    @property
+    def available(self) -> bool:
+        """Return True if entity is available."""
+        return super().available and self._lock.is_connected
+
 
 class TedeeDescriptionEntity(TedeeEntity):
     """Base class for Tedee device entities."""

homeassistant/components/tedee/lock.py

@@ -74,11 +74,6 @@ class TedeeLockEntity(TedeeEntity, LockEntity):
         """Return true if lock is jammed."""
         return self._lock.is_state_jammed
 
-    @property
-    def available(self) -> bool:
-        """Return True if entity is available."""
-        return super().available and self._lock.is_connected
-
     async def async_unlock(self, **kwargs: Any) -> None:
         """Unlock the door."""
         try:

tests/components/tedee/test_lock.py

@@ -25,7 +25,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 
-from tests.common import async_fire_time_changed
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 pytestmark = pytest.mark.usefixtures("init_integration")
 
@@ -210,6 +210,36 @@ async def test_update_failed(
     assert state.state == STATE_UNAVAILABLE
 
 
+async def test_cleanup_removed_locks(
+    hass: HomeAssistant,
+    mock_tedee: MagicMock,
+    device_registry: dr.DeviceRegistry,
+    mock_config_entry: MockConfigEntry,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Ensure removed locks are cleaned up."""
+
+    devices = dr.async_entries_for_config_entry(
+        device_registry, mock_config_entry.entry_id
+    )
+
+    locks = [device.name for device in devices]
+    assert "Lock-1A2B" in locks
+
+    # remove a lock and wait for coordinator
+    mock_tedee.locks_dict.pop(12345)
+    freezer.tick(timedelta(minutes=10))
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    devices = dr.async_entries_for_config_entry(
+        device_registry, mock_config_entry.entry_id
+    )
+
+    locks = [device.name for device in devices]
+    assert "Lock-1A2B" not in locks
+
+
 async def test_new_lock(
     hass: HomeAssistant,
     mock_tedee: MagicMock,