From f6850c36fcc5e6dadb938abce782df942dfb8e82 Mon Sep 17 00:00:00 2001
From: David Knowles <dknowles2@gmail.com>
Date: Sun, 6 Oct 2024 01:42:39 -0400
Subject: [PATCH] Fix problems with automatic management of Schlage locks
 (#127689)

Use the correct identifiers for existing lock devices
---
 .../components/schlage/coordinator.py         | 14 +++++++++---
 tests/components/schlage/test_init.py         | 22 ++++++++++++++-----
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/homeassistant/components/schlage/coordinator.py b/homeassistant/components/schlage/coordinator.py
index 365fabb8ac7..53bb43751a9 100644
--- a/homeassistant/components/schlage/coordinator.py
+++ b/homeassistant/components/schlage/coordinator.py
@@ -90,13 +90,21 @@ class SchlageDataUpdateCoordinator(DataUpdateCoordinator[SchlageData]):
         devices = dr.async_entries_for_config_entry(
             device_registry, self.config_entry.entry_id
         )
-        previous_locks = {device.id for device in devices}
+        previous_locks = set()
+        previous_locks_by_lock_id = {}
+        for device in devices:
+            for domain, identifier in device.identifiers:
+                if domain == DOMAIN:
+                    previous_locks.add(identifier)
+                    previous_locks_by_lock_id[identifier] = device
+                    continue
         current_locks = set(self.data.locks.keys())
+
         if removed_locks := previous_locks - current_locks:
             LOGGER.debug("Removed locks: %s", ", ".join(removed_locks))
-            for device_id in removed_locks:
+            for lock_id in removed_locks:
                 device_registry.async_update_device(
-                    device_id=device_id,
+                    device_id=previous_locks_by_lock_id[lock_id].id,
                     remove_config_entry_id=self.config_entry.entry_id,
                 )
 
diff --git a/tests/components/schlage/test_init.py b/tests/components/schlage/test_init.py
index 1f18bdde218..e40fc83a7ac 100644
--- a/tests/components/schlage/test_init.py
+++ b/tests/components/schlage/test_init.py
@@ -12,6 +12,7 @@ from syrupy.assertion import SnapshotAssertion
 from homeassistant.components.schlage.const import DOMAIN, UPDATE_INTERVAL
 from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.core import HomeAssistant
+import homeassistant.helpers.device_registry as dr
 from homeassistant.helpers.device_registry import DeviceRegistry
 
 from tests.common import MockConfigEntry, async_fire_time_changed
@@ -125,6 +126,10 @@ async def test_auto_add_device(
     """Test new devices are auto-added to the device registry."""
     device = device_registry.async_get_device(identifiers={(DOMAIN, "test")})
     assert device is not None
+    all_devices = dr.async_entries_for_config_entry(
+        device_registry, mock_added_config_entry.entry_id
+    )
+    assert len(all_devices) == 1
 
     mock_lock_attrs["device_id"] = "test2"
     new_mock_lock = create_autospec(Lock)
@@ -139,19 +144,21 @@ async def test_auto_add_device(
     new_device = device_registry.async_get_device(identifiers={(DOMAIN, "test2")})
     assert new_device is not None
 
+    all_devices = dr.async_entries_for_config_entry(
+        device_registry, mock_added_config_entry.entry_id
+    )
+    assert len(all_devices) == 2
+
 
 async def test_auto_remove_device(
     hass: HomeAssistant,
     device_registry: DeviceRegistry,
     mock_added_config_entry: ConfigEntry,
     mock_schlage: Mock,
-    mock_lock: Mock,
-    mock_lock_attrs: dict[str, Any],
     freezer: FrozenDateTimeFactory,
 ) -> None:
     """Test new devices are auto-added to the device registry."""
-    device = device_registry.async_get_device(identifiers={(DOMAIN, "test")})
-    assert device is not None
+    assert device_registry.async_get_device(identifiers={(DOMAIN, "test")}) is not None
 
     mock_schlage.locks.return_value = []
 
@@ -160,5 +167,8 @@ async def test_auto_remove_device(
     async_fire_time_changed(hass)
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    new_device = device_registry.async_get_device(identifiers={(DOMAIN, "test")})
-    assert new_device is None
+    assert device_registry.async_get_device(identifiers={(DOMAIN, "test")}) is None
+    all_devices = dr.async_entries_for_config_entry(
+        device_registry, mock_added_config_entry.entry_id
+    )
+    assert len(all_devices) == 0