Automatically add and remove Schlage devices (#125520)

* Allow manual deletion of stale Schlage devices

* Automatically add and remove locks

* Add tests and fix discovered bugs

* Changes requested during review

homeassistant/components/schlage/binary_sensor.py

@@ -45,15 +45,20 @@ async def async_setup_entry(
 ) -> None:
     """Set up binary_sensors based on a config entry."""
     coordinator: SchlageDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
-    async_add_entities(
+
-        SchlageBinarySensor(
+    def _add_new_locks(locks: dict[str, LockData]) -> None:
-            coordinator=coordinator,
+        async_add_entities(
-            description=description,
+            SchlageBinarySensor(
-            device_id=device_id,
+                coordinator=coordinator,
+                description=description,
+                device_id=device_id,
+            )
+            for device_id in locks
+            for description in _DESCRIPTIONS
         )
-        for device_id in coordinator.data.locks
+
-        for description in _DESCRIPTIONS
+    _add_new_locks(coordinator.data.locks)
-    )
+    coordinator.new_locks_callbacks.append(_add_new_locks)
 
 
 class SchlageBinarySensor(SchlageEntity, BinarySensorEntity):

homeassistant/components/schlage/coordinator.py

@@ -3,14 +3,17 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Callable
 from dataclasses import dataclass
 
 from pyschlage import Lock, Schlage
 from pyschlage.exceptions import Error as SchlageError, NotAuthorizedError
 from pyschlage.log import LockLog
 
-from homeassistant.core import HomeAssistant
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryAuthFailed
+import homeassistant.helpers.device_registry as dr
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import DOMAIN, LOGGER, UPDATE_INTERVAL
@@ -34,12 +37,16 @@ class SchlageData:
 class SchlageDataUpdateCoordinator(DataUpdateCoordinator[SchlageData]):
     """The Schlage data update coordinator."""
 
+    config_entry: ConfigEntry
+
     def __init__(self, hass: HomeAssistant, username: str, api: Schlage) -> None:
         """Initialize the class."""
         super().__init__(
             hass, LOGGER, name=f"{DOMAIN} ({username})", update_interval=UPDATE_INTERVAL
         )
         self.api = api
+        self.new_locks_callbacks: list[Callable[[dict[str, LockData]], None]] = []
+        self.async_add_listener(self._add_remove_locks)
 
     async def _async_update_data(self) -> SchlageData:
         """Fetch the latest data from the Schlage API."""
@@ -55,9 +62,7 @@ class SchlageDataUpdateCoordinator(DataUpdateCoordinator[SchlageData]):
                 for lock in locks
             )
         )
-        return SchlageData(
+        return SchlageData(locks={ld.lock.device_id: ld for ld in lock_data})
-            locks={ld.lock.device_id: ld for ld in lock_data},
-        )
 
     def _get_lock_data(self, lock: Lock) -> LockData:
         logs: list[LockLog] = []
@@ -74,3 +79,29 @@ class SchlageDataUpdateCoordinator(DataUpdateCoordinator[SchlageData]):
             LOGGER.debug('Failed to read logs for lock "%s": %s', lock.name, ex)
 
         return LockData(lock=lock, logs=logs)
+
+    @callback
+    def _add_remove_locks(self) -> None:
+        """Add newly discovered locks and remove nonexistent locks."""
+        if self.data is None:
+            return
+
+        device_registry = dr.async_get(self.hass)
+        devices = dr.async_entries_for_config_entry(
+            device_registry, self.config_entry.entry_id
+        )
+        previous_locks = {device.id for device in devices}
+        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:
+                device_registry.async_update_device(
+                    device_id=device_id,
+                    remove_config_entry_id=self.config_entry.entry_id,
+                )
+
+        if new_lock_ids := current_locks - previous_locks:
+            LOGGER.debug("New locks found: %s", ", ".join(new_lock_ids))
+            new_locks = {lock_id: self.data.locks[lock_id] for lock_id in new_lock_ids}
+            for new_lock_callback in self.new_locks_callbacks:
+                new_lock_callback(new_locks)

homeassistant/components/schlage/lock.py

@@ -10,7 +10,7 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .coordinator import SchlageDataUpdateCoordinator
+from .coordinator import LockData, SchlageDataUpdateCoordinator
 from .entity import SchlageEntity
 
 
@@ -21,10 +21,15 @@ async def async_setup_entry(
 ) -> None:
     """Set up Schlage WiFi locks based on a config entry."""
     coordinator: SchlageDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
-    async_add_entities(
+
-        SchlageLockEntity(coordinator=coordinator, device_id=device_id)
+    def _add_new_locks(locks: dict[str, LockData]) -> None:
-        for device_id in coordinator.data.locks
+        async_add_entities(
-    )
+            SchlageLockEntity(coordinator=coordinator, device_id=device_id)
+            for device_id in locks
+        )
+
+    _add_new_locks(coordinator.data.locks)
+    coordinator.new_locks_callbacks.append(_add_new_locks)
 
 
 class SchlageLockEntity(SchlageEntity, LockEntity):

homeassistant/components/schlage/sensor.py

@@ -14,7 +14,7 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .coordinator import SchlageDataUpdateCoordinator
+from .coordinator import LockData, SchlageDataUpdateCoordinator
 from .entity import SchlageEntity
 
 _SENSOR_DESCRIPTIONS: list[SensorEntityDescription] = [
@@ -35,15 +35,20 @@ async def async_setup_entry(
 ) -> None:
     """Set up sensors based on a config entry."""
     coordinator: SchlageDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
-    async_add_entities(
+
-        SchlageBatterySensor(
+    def _add_new_locks(locks: dict[str, LockData]) -> None:
-            coordinator=coordinator,
+        async_add_entities(
-            description=description,
+            SchlageBatterySensor(
-            device_id=device_id,
+                coordinator=coordinator,
+                description=description,
+                device_id=device_id,
+            )
+            for description in _SENSOR_DESCRIPTIONS
+            for device_id in locks
         )
-        for description in _SENSOR_DESCRIPTIONS
+
-        for device_id in coordinator.data.locks
+    _add_new_locks(coordinator.data.locks)
-    )
+    coordinator.new_locks_callbacks.append(_add_new_locks)
 
 
 class SchlageBatterySensor(SchlageEntity, SensorEntity):

homeassistant/components/schlage/switch.py

@@ -20,7 +20,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .coordinator import SchlageDataUpdateCoordinator
+from .coordinator import LockData, SchlageDataUpdateCoordinator
 from .entity import SchlageEntity
 
 
@@ -62,15 +62,20 @@ async def async_setup_entry(
 ) -> None:
     """Set up switches based on a config entry."""
     coordinator: SchlageDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
-    async_add_entities(
+
-        SchlageSwitch(
+    def _add_new_locks(locks: dict[str, LockData]) -> None:
-            coordinator=coordinator,
+        async_add_entities(
-            description=description,
+            SchlageSwitch(
-            device_id=device_id,
+                coordinator=coordinator,
+                description=description,
+                device_id=device_id,
+            )
+            for device_id in locks
+            for description in SWITCHES
         )
-        for device_id in coordinator.data.locks
+
-        for description in SWITCHES
+    _add_new_locks(coordinator.data.locks)
-    )
+    coordinator.new_locks_callbacks.append(_add_new_locks)
 
 
 class SchlageSwitch(SchlageEntity, SwitchEntity):

tests/components/schlage/conftest.py

@@ -1,6 +1,7 @@
 """Common fixtures for the Schlage tests."""
 
 from collections.abc import Generator
+from typing import Any
 from unittest.mock import AsyncMock, Mock, create_autospec, patch
 
 from pyschlage.lock import Lock
@@ -70,21 +71,27 @@ def mock_pyschlage_auth() -> Mock:
 
 
 @pytest.fixture
-def mock_lock() -> Mock:
+def mock_lock(mock_lock_attrs: dict[str, Any]) -> Mock:
     """Mock Lock fixture."""
     mock_lock = create_autospec(Lock)
-    mock_lock.configure_mock(
+    mock_lock.configure_mock(**mock_lock_attrs)
-        device_id="test",
-        name="Vault Door",
-        model_name="<model-name>",
-        is_locked=False,
-        is_jammed=False,
-        battery_level=20,
-        firmware_version="1.0",
-        lock_and_leave_enabled=True,
-        beeper_enabled=True,
-    )
     mock_lock.logs.return_value = []
     mock_lock.last_changed_by.return_value = "thumbturn"
     mock_lock.keypad_disabled.return_value = False
     return mock_lock
+
+
+@pytest.fixture
+def mock_lock_attrs() -> dict[str, Any]:
+    """Attributes for a mock lock."""
+    return {
+        "device_id": "test",
+        "name": "Vault Door",
+        "model_name": "<model-name>",
+        "is_locked": False,
+        "is_jammed": False,
+        "battery_level": 20,
+        "firmware_version": "1.0",
+        "lock_and_leave_enabled": True,
+        "beeper_enabled": True,
+    }

tests/components/schlage/snapshots/test_init.ambr

@@ -0,0 +1,33 @@
+# serializer version: 1
+# name: test_lock_device_registry
+  DeviceRegistryEntrySnapshot({
+    'area_id': None,
+    'config_entries': <ANY>,
+    'configuration_url': None,
+    'connections': set({
+    }),
+    'disabled_by': None,
+    'entry_type': None,
+    'hw_version': None,
+    'id': <ANY>,
+    'identifiers': set({
+      tuple(
+        'schlage',
+        'test',
+      ),
+    }),
+    'is_new': False,
+    'labels': set({
+    }),
+    'manufacturer': 'Schlage',
+    'model': '<model-name>',
+    'model_id': None,
+    'name': 'Vault Door',
+    'name_by_user': None,
+    'primary_config_entry': <ANY>,
+    'serial_number': None,
+    'suggested_area': None,
+    'sw_version': '1.0',
+    'via_device_id': None,
+  })
+# ---

tests/components/schlage/test_binary_sensor.py

@@ -8,7 +8,7 @@ from pyschlage.exceptions import UnknownError
 
 from homeassistant.components.binary_sensor import BinarySensorDeviceClass
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import STATE_ON, STATE_UNAVAILABLE
+from homeassistant.const import STATE_ON
 from homeassistant.core import HomeAssistant
 
 from tests.common import async_fire_time_changed
@@ -37,15 +37,6 @@ async def test_keypad_disabled_binary_sensor(
 
     mock_lock.keypad_disabled.assert_called_once_with([])
 
-    mock_schlage.locks.return_value = []
-    # Make the coordinator refresh data.
-    freezer.tick(timedelta(seconds=30))
-    async_fire_time_changed(hass)
-    await hass.async_block_till_done(wait_background_tasks=True)
-    keypad = hass.states.get("binary_sensor.vault_door_keypad_disabled")
-    assert keypad is not None
-    assert keypad.state == STATE_UNAVAILABLE
-
 
 async def test_keypad_disabled_binary_sensor_use_previous_logs_on_failure(
     hass: HomeAssistant,

tests/components/schlage/test_init.py

@@ -1,14 +1,20 @@
 """Tests for the Schlage integration."""
 
-from unittest.mock import Mock, patch
+from typing import Any
+from unittest.mock import Mock, create_autospec, patch
 
+from freezegun.api import FrozenDateTimeFactory
 from pycognito.exceptions import WarrantException
 from pyschlage.exceptions import Error, NotAuthorizedError
+from pyschlage.lock import Lock
+from syrupy.assertion import SnapshotAssertion
 
-from homeassistant.config_entries import ConfigEntryState
+from homeassistant.components.schlage.const import DOMAIN, UPDATE_INTERVAL
+from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.device_registry import DeviceRegistry
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 
 @patch(
@@ -94,3 +100,65 @@ async def test_load_unload_config_entry(
     await hass.config_entries.async_unload(mock_config_entry.entry_id)
     await hass.async_block_till_done()
     assert mock_config_entry.state is ConfigEntryState.NOT_LOADED
+
+
+async def test_lock_device_registry(
+    hass: HomeAssistant,
+    device_registry: DeviceRegistry,
+    mock_added_config_entry: ConfigEntry,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test lock is added to device registry."""
+    device = device_registry.async_get_device(identifiers={(DOMAIN, "test")})
+    assert device == snapshot
+
+
+async def test_auto_add_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
+
+    mock_lock_attrs["device_id"] = "test2"
+    new_mock_lock = create_autospec(Lock)
+    new_mock_lock.configure_mock(**mock_lock_attrs)
+    mock_schlage.locks.return_value = [mock_lock, new_mock_lock]
+
+    # Make the coordinator refresh data.
+    freezer.tick(UPDATE_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    new_device = device_registry.async_get_device(identifiers={(DOMAIN, "test2")})
+    assert new_device is not None
+
+
+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
+
+    mock_schlage.locks.return_value = []
+
+    # Make the coordinator refresh data.
+    freezer.tick(UPDATE_INTERVAL)
+    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

tests/components/schlage/test_lock.py

@@ -12,28 +12,13 @@ from homeassistant.const import (
     SERVICE_LOCK,
     SERVICE_UNLOCK,
     STATE_JAMMED,
-    STATE_UNAVAILABLE,
     STATE_UNLOCKED,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry as dr
 
 from tests.common import async_fire_time_changed
 
 
-async def test_lock_device_registry(
-    hass: HomeAssistant,
-    device_registry: dr.DeviceRegistry,
-    mock_added_config_entry: ConfigEntry,
-) -> None:
-    """Test lock is added to device registry."""
-    device = device_registry.async_get_device(identifiers={("schlage", "test")})
-    assert device.model == "<model-name>"
-    assert device.sw_version == "1.0"
-    assert device.name == "Vault Door"
-    assert device.manufacturer == "Schlage"
-
-
 async def test_lock_attributes(
     hass: HomeAssistant,
     mock_added_config_entry: ConfigEntry,
@@ -57,16 +42,6 @@ async def test_lock_attributes(
     assert lock is not None
     assert lock.state == STATE_JAMMED
 
-    mock_schlage.locks.return_value = []
-    # Make the coordinator refresh data.
-    freezer.tick(timedelta(seconds=30))
-    async_fire_time_changed(hass)
-    await hass.async_block_till_done(wait_background_tasks=True)
-    lock = hass.states.get("lock.vault_door")
-    assert lock is not None
-    assert lock.state == STATE_UNAVAILABLE
-    assert "changed_by" not in lock.attributes
-
 
 async def test_lock_services(
     hass: HomeAssistant, mock_lock: Mock, mock_added_config_entry: ConfigEntry
@@ -107,7 +82,7 @@ async def test_changed_by(
     freezer.tick(timedelta(seconds=30))
     async_fire_time_changed(hass)
     await hass.async_block_till_done(wait_background_tasks=True)
-    mock_lock.last_changed_by.assert_called_once_with()
+    mock_lock.last_changed_by.assert_called_with()
 
     lock_device = hass.states.get("lock.vault_door")
     assert lock_device is not None

tests/components/schlage/test_sensor.py

@@ -4,20 +4,6 @@ from homeassistant.components.sensor import SensorDeviceClass
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry as dr
-
-
-async def test_sensor_device_registry(
-    hass: HomeAssistant,
-    device_registry: dr.DeviceRegistry,
-    mock_added_config_entry: ConfigEntry,
-) -> None:
-    """Test sensor is added to device registry."""
-    device = device_registry.async_get_device(identifiers={("schlage", "test")})
-    assert device.model == "<model-name>"
-    assert device.sw_version == "1.0"
-    assert device.name == "Vault Door"
-    assert device.manufacturer == "Schlage"
 
 
 async def test_battery_sensor(

tests/components/schlage/test_switch.py

@@ -6,20 +6,6 @@ from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_ENTITY_ID, SERVICE_TURN_OFF, SERVICE_TURN_ON
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry as dr
-
-
-async def test_switch_device_registry(
-    hass: HomeAssistant,
-    device_registry: dr.DeviceRegistry,
-    mock_added_config_entry: ConfigEntry,
-) -> None:
-    """Test switch is added to device registry."""
-    device = device_registry.async_get_device(identifiers={("schlage", "test")})
-    assert device.model == "<model-name>"
-    assert device.sw_version == "1.0"
-    assert device.name == "Vault Door"
-    assert device.manufacturer == "Schlage"
 
 
 async def test_beeper_services(