From 0af913cc9a9d8f200f30b05577bc6d9c03f5429e Mon Sep 17 00:00:00 2001 From: David Knowles Date: Fri, 13 Sep 2024 09:17:51 -0400 Subject: [PATCH] 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 --- .../components/schlage/binary_sensor.py | 21 ++++-- .../components/schlage/coordinator.py | 39 +++++++++- homeassistant/components/schlage/lock.py | 15 ++-- homeassistant/components/schlage/sensor.py | 23 +++--- homeassistant/components/schlage/switch.py | 23 +++--- tests/components/schlage/conftest.py | 31 +++++--- .../schlage/snapshots/test_init.ambr | 33 +++++++++ .../components/schlage/test_binary_sensor.py | 11 +-- tests/components/schlage/test_init.py | 74 ++++++++++++++++++- tests/components/schlage/test_lock.py | 27 +------ tests/components/schlage/test_sensor.py | 14 ---- tests/components/schlage/test_switch.py | 14 ---- 12 files changed, 211 insertions(+), 114 deletions(-) create mode 100644 tests/components/schlage/snapshots/test_init.ambr diff --git a/homeassistant/components/schlage/binary_sensor.py b/homeassistant/components/schlage/binary_sensor.py index a141403bdf4..bc1ee666f9e 100644 --- a/homeassistant/components/schlage/binary_sensor.py +++ b/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( - coordinator=coordinator, - description=description, - device_id=device_id, + + def _add_new_locks(locks: dict[str, LockData]) -> None: + async_add_entities( + SchlageBinarySensor( + 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): diff --git a/homeassistant/components/schlage/coordinator.py b/homeassistant/components/schlage/coordinator.py index 959d1e215f8..365fabb8ac7 100644 --- a/homeassistant/components/schlage/coordinator.py +++ b/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( - locks={ld.lock.device_id: ld for ld in lock_data}, - ) + return SchlageData(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) diff --git a/homeassistant/components/schlage/lock.py b/homeassistant/components/schlage/lock.py index 59ce00e809a..97dbfc78d41 100644 --- a/homeassistant/components/schlage/lock.py +++ b/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) - for device_id in coordinator.data.locks - ) + + def _add_new_locks(locks: dict[str, LockData]) -> None: + 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): diff --git a/homeassistant/components/schlage/sensor.py b/homeassistant/components/schlage/sensor.py index 8de09fa4cbb..115412882a2 100644 --- a/homeassistant/components/schlage/sensor.py +++ b/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( - coordinator=coordinator, - description=description, - device_id=device_id, + + def _add_new_locks(locks: dict[str, LockData]) -> None: + async_add_entities( + SchlageBatterySensor( + 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): diff --git a/homeassistant/components/schlage/switch.py b/homeassistant/components/schlage/switch.py index 53771768ccd..aaed57fc741 100644 --- a/homeassistant/components/schlage/switch.py +++ b/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( - coordinator=coordinator, - description=description, - device_id=device_id, + + def _add_new_locks(locks: dict[str, LockData]) -> None: + async_add_entities( + SchlageSwitch( + 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): diff --git a/tests/components/schlage/conftest.py b/tests/components/schlage/conftest.py index 9d61bb877d9..5ff8d045606 100644 --- a/tests/components/schlage/conftest.py +++ b/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( - device_id="test", - name="Vault Door", - 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.configure_mock(**mock_lock_attrs) 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": "", + "is_locked": False, + "is_jammed": False, + "battery_level": 20, + "firmware_version": "1.0", + "lock_and_leave_enabled": True, + "beeper_enabled": True, + } diff --git a/tests/components/schlage/snapshots/test_init.ambr b/tests/components/schlage/snapshots/test_init.ambr new file mode 100644 index 00000000000..c7049443ab7 --- /dev/null +++ b/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': , + 'configuration_url': None, + 'connections': set({ + }), + 'disabled_by': None, + 'entry_type': None, + 'hw_version': None, + 'id': , + 'identifiers': set({ + tuple( + 'schlage', + 'test', + ), + }), + 'is_new': False, + 'labels': set({ + }), + 'manufacturer': 'Schlage', + 'model': '', + 'model_id': None, + 'name': 'Vault Door', + 'name_by_user': None, + 'primary_config_entry': , + 'serial_number': None, + 'suggested_area': None, + 'sw_version': '1.0', + 'via_device_id': None, + }) +# --- diff --git a/tests/components/schlage/test_binary_sensor.py b/tests/components/schlage/test_binary_sensor.py index dbbc5b07b87..91bd996ba5b 100644 --- a/tests/components/schlage/test_binary_sensor.py +++ b/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, diff --git a/tests/components/schlage/test_init.py b/tests/components/schlage/test_init.py index 0fe7af1982b..1f18bdde218 100644 --- a/tests/components/schlage/test_init.py +++ b/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 diff --git a/tests/components/schlage/test_lock.py b/tests/components/schlage/test_lock.py index ab0f4f5d863..74af80dce84 100644 --- a/tests/components/schlage/test_lock.py +++ b/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 == "" - 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 diff --git a/tests/components/schlage/test_sensor.py b/tests/components/schlage/test_sensor.py index 2c0cabbb1e8..9fa90edecbb 100644 --- a/tests/components/schlage/test_sensor.py +++ b/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 == "" - assert device.sw_version == "1.0" - assert device.name == "Vault Door" - assert device.manufacturer == "Schlage" async def test_battery_sensor( diff --git a/tests/components/schlage/test_switch.py b/tests/components/schlage/test_switch.py index f1cded3ce22..52b8da81670 100644 --- a/tests/components/schlage/test_switch.py +++ b/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 == "" - assert device.sw_version == "1.0" - assert device.name == "Vault Door" - assert device.manufacturer == "Schlage" async def test_beeper_services(