From a0ea9a1e83b496f1028aaf60d7bd4263db5a513f Mon Sep 17 00:00:00 2001 From: David Knowles Date: Thu, 28 Nov 2024 02:29:29 -0500 Subject: [PATCH] Store Schlage runtime data in entry.runtime_data (#131731) --- homeassistant/components/schlage/__init__.py | 14 +++++------- .../components/schlage/binary_sensor.py | 7 +++--- .../components/schlage/diagnostics.py | 8 +++---- homeassistant/components/schlage/lock.py | 7 +++--- homeassistant/components/schlage/select.py | 7 +++--- homeassistant/components/schlage/sensor.py | 3 +-- homeassistant/components/schlage/switch.py | 3 +-- tests/components/schlage/__init__.py | 6 +++++ tests/components/schlage/conftest.py | 8 ++++--- .../components/schlage/test_binary_sensor.py | 7 +++--- tests/components/schlage/test_config_flow.py | 13 +++++------ tests/components/schlage/test_diagnostics.py | 5 +++-- tests/components/schlage/test_init.py | 22 ++++++++++--------- tests/components/schlage/test_lock.py | 11 ++++++---- tests/components/schlage/test_select.py | 7 ++++-- tests/components/schlage/test_sensor.py | 5 +++-- tests/components/schlage/test_switch.py | 11 +++++++--- 17 files changed, 78 insertions(+), 66 deletions(-) diff --git a/homeassistant/components/schlage/__init__.py b/homeassistant/components/schlage/__init__.py index e9fb24f1309..6eae69d9542 100644 --- a/homeassistant/components/schlage/__init__.py +++ b/homeassistant/components/schlage/__init__.py @@ -10,7 +10,6 @@ from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed -from .const import DOMAIN from .coordinator import SchlageDataUpdateCoordinator PLATFORMS: list[Platform] = [ @@ -21,8 +20,10 @@ PLATFORMS: list[Platform] = [ Platform.SWITCH, ] +type SchlageConfigEntry = ConfigEntry[SchlageDataUpdateCoordinator] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + +async def async_setup_entry(hass: HomeAssistant, entry: SchlageConfigEntry) -> bool: """Set up Schlage from a config entry.""" username = entry.data[CONF_USERNAME] password = entry.data[CONF_PASSWORD] @@ -32,15 +33,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: raise ConfigEntryAuthFailed from ex coordinator = SchlageDataUpdateCoordinator(hass, username, pyschlage.Schlage(auth)) - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + entry.runtime_data = coordinator await coordinator.async_config_entry_first_refresh() await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: SchlageConfigEntry) -> bool: """Unload a config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - hass.data[DOMAIN].pop(entry.entry_id) - - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/schlage/binary_sensor.py b/homeassistant/components/schlage/binary_sensor.py index bc1ee666f9e..f928d42b3ee 100644 --- a/homeassistant/components/schlage/binary_sensor.py +++ b/homeassistant/components/schlage/binary_sensor.py @@ -10,12 +10,11 @@ from homeassistant.components.binary_sensor import ( BinarySensorEntity, BinarySensorEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import SchlageConfigEntry from .coordinator import LockData, SchlageDataUpdateCoordinator from .entity import SchlageEntity @@ -40,11 +39,11 @@ _DESCRIPTIONS: tuple[SchlageBinarySensorEntityDescription] = ( async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: SchlageConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up binary_sensors based on a config entry.""" - coordinator: SchlageDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data def _add_new_locks(locks: dict[str, LockData]) -> None: async_add_entities( diff --git a/homeassistant/components/schlage/diagnostics.py b/homeassistant/components/schlage/diagnostics.py index af1bf311676..ec4d9c489e3 100644 --- a/homeassistant/components/schlage/diagnostics.py +++ b/homeassistant/components/schlage/diagnostics.py @@ -4,19 +4,17 @@ from __future__ import annotations from typing import Any -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from .const import DOMAIN -from .coordinator import SchlageDataUpdateCoordinator +from . import SchlageConfigEntry async def async_get_config_entry_diagnostics( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: SchlageConfigEntry, ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - coordinator: SchlageDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data # NOTE: Schlage diagnostics are already redacted. return { "locks": [ld.lock.get_diagnostics() for ld in coordinator.data.locks.values()] diff --git a/homeassistant/components/schlage/lock.py b/homeassistant/components/schlage/lock.py index 97dbfc78d41..d203913191d 100644 --- a/homeassistant/components/schlage/lock.py +++ b/homeassistant/components/schlage/lock.py @@ -5,22 +5,21 @@ from __future__ import annotations from typing import Any from homeassistant.components.lock import LockEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import SchlageConfigEntry from .coordinator import LockData, SchlageDataUpdateCoordinator from .entity import SchlageEntity async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: SchlageConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Schlage WiFi locks based on a config entry.""" - coordinator: SchlageDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data def _add_new_locks(locks: dict[str, LockData]) -> None: async_add_entities( diff --git a/homeassistant/components/schlage/select.py b/homeassistant/components/schlage/select.py index 6d93eccaa85..6cf0853835f 100644 --- a/homeassistant/components/schlage/select.py +++ b/homeassistant/components/schlage/select.py @@ -3,12 +3,11 @@ from __future__ import annotations from homeassistant.components.select import SelectEntity, SelectEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import SchlageConfigEntry from .coordinator import LockData, SchlageDataUpdateCoordinator from .entity import SchlageEntity @@ -33,11 +32,11 @@ _DESCRIPTIONS = ( async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: SchlageConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up selects based on a config entry.""" - coordinator: SchlageDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data def _add_new_locks(locks: dict[str, LockData]) -> None: async_add_entities( diff --git a/homeassistant/components/schlage/sensor.py b/homeassistant/components/schlage/sensor.py index 115412882a2..a15d1740b91 100644 --- a/homeassistant/components/schlage/sensor.py +++ b/homeassistant/components/schlage/sensor.py @@ -13,7 +13,6 @@ from homeassistant.const import PERCENTAGE, EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN from .coordinator import LockData, SchlageDataUpdateCoordinator from .entity import SchlageEntity @@ -34,7 +33,7 @@ async def async_setup_entry( async_add_entities: AddEntitiesCallback, ) -> None: """Set up sensors based on a config entry.""" - coordinator: SchlageDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data def _add_new_locks(locks: dict[str, LockData]) -> None: async_add_entities( diff --git a/homeassistant/components/schlage/switch.py b/homeassistant/components/schlage/switch.py index aaed57fc741..39fe6dbbc99 100644 --- a/homeassistant/components/schlage/switch.py +++ b/homeassistant/components/schlage/switch.py @@ -19,7 +19,6 @@ from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN from .coordinator import LockData, SchlageDataUpdateCoordinator from .entity import SchlageEntity @@ -61,7 +60,7 @@ async def async_setup_entry( async_add_entities: AddEntitiesCallback, ) -> None: """Set up switches based on a config entry.""" - coordinator: SchlageDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data def _add_new_locks(locks: dict[str, LockData]) -> None: async_add_entities( diff --git a/tests/components/schlage/__init__.py b/tests/components/schlage/__init__.py index c6cd3fec0bc..613621b2fb8 100644 --- a/tests/components/schlage/__init__.py +++ b/tests/components/schlage/__init__.py @@ -1 +1,7 @@ """Tests for the Schlage integration.""" + +from homeassistant.components.schlage.coordinator import SchlageDataUpdateCoordinator + +from tests.common import MockConfigEntry + +type MockSchlageConfigEntry = MockConfigEntry[SchlageDataUpdateCoordinator] diff --git a/tests/components/schlage/conftest.py b/tests/components/schlage/conftest.py index f774b8cfb89..6695191dcf0 100644 --- a/tests/components/schlage/conftest.py +++ b/tests/components/schlage/conftest.py @@ -11,11 +11,13 @@ from homeassistant.components.schlage.const import DOMAIN from homeassistant.const import CONF_PASSWORD, CONF_USERNAME from homeassistant.core import HomeAssistant +from . import MockSchlageConfigEntry + from tests.common import MockConfigEntry @pytest.fixture -def mock_config_entry() -> MockConfigEntry: +def mock_config_entry() -> MockSchlageConfigEntry: """Mock ConfigEntry.""" return MockConfigEntry( title="asdf@asdf.com", @@ -31,11 +33,11 @@ def mock_config_entry() -> MockConfigEntry: @pytest.fixture async def mock_added_config_entry( hass: HomeAssistant, - mock_config_entry: MockConfigEntry, + mock_config_entry: MockSchlageConfigEntry, mock_pyschlage_auth: Mock, mock_schlage: Mock, mock_lock: Mock, -) -> MockConfigEntry: +) -> MockSchlageConfigEntry: """Mock ConfigEntry that's been added to HA.""" mock_schlage.locks.return_value = [mock_lock] mock_schlage.users.return_value = [] diff --git a/tests/components/schlage/test_binary_sensor.py b/tests/components/schlage/test_binary_sensor.py index 91bd996ba5b..a073097f755 100644 --- a/tests/components/schlage/test_binary_sensor.py +++ b/tests/components/schlage/test_binary_sensor.py @@ -7,10 +7,11 @@ from freezegun.api import FrozenDateTimeFactory from pyschlage.exceptions import UnknownError from homeassistant.components.binary_sensor import BinarySensorDeviceClass -from homeassistant.config_entries import ConfigEntry from homeassistant.const import STATE_ON from homeassistant.core import HomeAssistant +from . import MockSchlageConfigEntry + from tests.common import async_fire_time_changed @@ -18,7 +19,7 @@ async def test_keypad_disabled_binary_sensor( hass: HomeAssistant, mock_schlage: Mock, mock_lock: Mock, - mock_added_config_entry: ConfigEntry, + mock_added_config_entry: MockSchlageConfigEntry, freezer: FrozenDateTimeFactory, ) -> None: """Test the keypad_disabled binary_sensor.""" @@ -42,7 +43,7 @@ async def test_keypad_disabled_binary_sensor_use_previous_logs_on_failure( hass: HomeAssistant, mock_schlage: Mock, mock_lock: Mock, - mock_added_config_entry: ConfigEntry, + mock_added_config_entry: MockSchlageConfigEntry, freezer: FrozenDateTimeFactory, ) -> None: """Test the keypad_disabled binary_sensor.""" diff --git a/tests/components/schlage/test_config_flow.py b/tests/components/schlage/test_config_flow.py index 7f4a40f9b53..88b5f113863 100644 --- a/tests/components/schlage/test_config_flow.py +++ b/tests/components/schlage/test_config_flow.py @@ -10,7 +10,7 @@ from homeassistant.components.schlage.const import DOMAIN from homeassistant.core import HomeAssistant from homeassistant.data_entry_flow import FlowResultType -from tests.common import MockConfigEntry +from . import MockSchlageConfigEntry pytestmark = pytest.mark.usefixtures("mock_setup_entry") @@ -95,8 +95,7 @@ async def test_form_unknown(hass: HomeAssistant, mock_pyschlage_auth: Mock) -> N async def test_reauth( hass: HomeAssistant, - mock_added_config_entry: MockConfigEntry, - mock_setup_entry: AsyncMock, + mock_added_config_entry: MockSchlageConfigEntry, mock_pyschlage_auth: Mock, ) -> None: """Test reauth flow.""" @@ -104,8 +103,7 @@ async def test_reauth( await hass.async_block_till_done() flows = hass.config_entries.flow.async_progress() - assert len(flows) == 1 - [result] = flows + result = flows[-1] assert result["step_id"] == "reauth_confirm" result2 = await hass.config_entries.flow.async_configure( @@ -121,12 +119,11 @@ async def test_reauth( "username": "asdf@asdf.com", "password": "new-password", } - assert len(mock_setup_entry.mock_calls) == 1 async def test_reauth_invalid_auth( hass: HomeAssistant, - mock_added_config_entry: MockConfigEntry, + mock_added_config_entry: MockSchlageConfigEntry, mock_setup_entry: AsyncMock, mock_pyschlage_auth: Mock, ) -> None: @@ -154,7 +151,7 @@ async def test_reauth_invalid_auth( async def test_reauth_wrong_account( hass: HomeAssistant, - mock_added_config_entry: MockConfigEntry, + mock_added_config_entry: MockSchlageConfigEntry, mock_setup_entry: AsyncMock, mock_pyschlage_auth: Mock, ) -> None: diff --git a/tests/components/schlage/test_diagnostics.py b/tests/components/schlage/test_diagnostics.py index 15b2316bf38..0b0dc856c1a 100644 --- a/tests/components/schlage/test_diagnostics.py +++ b/tests/components/schlage/test_diagnostics.py @@ -4,7 +4,8 @@ from unittest.mock import Mock from homeassistant.core import HomeAssistant -from tests.common import MockConfigEntry +from . import MockSchlageConfigEntry + from tests.components.diagnostics import get_diagnostics_for_config_entry from tests.typing import ClientSessionGenerator @@ -12,7 +13,7 @@ from tests.typing import ClientSessionGenerator async def test_entry_diagnostics( hass: HomeAssistant, hass_client: ClientSessionGenerator, - mock_added_config_entry: MockConfigEntry, + mock_added_config_entry: MockSchlageConfigEntry, mock_lock: Mock, ) -> None: """Test Schlage diagnostics.""" diff --git a/tests/components/schlage/test_init.py b/tests/components/schlage/test_init.py index e40fc83a7ac..57a139e582e 100644 --- a/tests/components/schlage/test_init.py +++ b/tests/components/schlage/test_init.py @@ -10,12 +10,14 @@ from pyschlage.lock import Lock from syrupy.assertion import SnapshotAssertion from homeassistant.components.schlage.const import DOMAIN, UPDATE_INTERVAL -from homeassistant.config_entries import ConfigEntry, ConfigEntryState +from homeassistant.config_entries import 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 +from . import MockSchlageConfigEntry + +from tests.common import async_fire_time_changed @patch( @@ -23,7 +25,7 @@ from tests.common import MockConfigEntry, async_fire_time_changed side_effect=WarrantException, ) async def test_auth_failed( - mock_auth: Mock, hass: HomeAssistant, mock_config_entry: MockConfigEntry + mock_auth: Mock, hass: HomeAssistant, mock_config_entry: MockSchlageConfigEntry ) -> None: """Test failed auth on setup.""" mock_config_entry.add_to_hass(hass) @@ -36,7 +38,7 @@ async def test_auth_failed( async def test_update_data_fails( hass: HomeAssistant, - mock_config_entry: MockConfigEntry, + mock_config_entry: MockSchlageConfigEntry, mock_pyschlage_auth: Mock, mock_schlage: Mock, ) -> None: @@ -52,7 +54,7 @@ async def test_update_data_fails( async def test_update_data_auth_error( hass: HomeAssistant, - mock_config_entry: MockConfigEntry, + mock_config_entry: MockSchlageConfigEntry, mock_pyschlage_auth: Mock, mock_schlage: Mock, ) -> None: @@ -68,7 +70,7 @@ async def test_update_data_auth_error( async def test_update_data_get_logs_auth_error( hass: HomeAssistant, - mock_config_entry: MockConfigEntry, + mock_config_entry: MockSchlageConfigEntry, mock_pyschlage_auth: Mock, mock_schlage: Mock, mock_lock: Mock, @@ -87,7 +89,7 @@ async def test_update_data_get_logs_auth_error( async def test_load_unload_config_entry( hass: HomeAssistant, - mock_config_entry: MockConfigEntry, + mock_config_entry: MockSchlageConfigEntry, mock_pyschlage_auth: Mock, mock_schlage: Mock, ) -> None: @@ -106,7 +108,7 @@ async def test_load_unload_config_entry( async def test_lock_device_registry( hass: HomeAssistant, device_registry: DeviceRegistry, - mock_added_config_entry: ConfigEntry, + mock_added_config_entry: MockSchlageConfigEntry, snapshot: SnapshotAssertion, ) -> None: """Test lock is added to device registry.""" @@ -117,7 +119,7 @@ async def test_lock_device_registry( async def test_auto_add_device( hass: HomeAssistant, device_registry: DeviceRegistry, - mock_added_config_entry: ConfigEntry, + mock_added_config_entry: MockSchlageConfigEntry, mock_schlage: Mock, mock_lock: Mock, mock_lock_attrs: dict[str, Any], @@ -153,7 +155,7 @@ async def test_auto_add_device( async def test_auto_remove_device( hass: HomeAssistant, device_registry: DeviceRegistry, - mock_added_config_entry: ConfigEntry, + mock_added_config_entry: MockSchlageConfigEntry, mock_schlage: Mock, freezer: FrozenDateTimeFactory, ) -> None: diff --git a/tests/components/schlage/test_lock.py b/tests/components/schlage/test_lock.py index 518c723d581..6a3bb799213 100644 --- a/tests/components/schlage/test_lock.py +++ b/tests/components/schlage/test_lock.py @@ -6,16 +6,17 @@ from unittest.mock import Mock from freezegun.api import FrozenDateTimeFactory from homeassistant.components.lock import DOMAIN as LOCK_DOMAIN, LockState -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ATTR_ENTITY_ID, SERVICE_LOCK, SERVICE_UNLOCK from homeassistant.core import HomeAssistant +from . import MockSchlageConfigEntry + from tests.common import async_fire_time_changed async def test_lock_attributes( hass: HomeAssistant, - mock_added_config_entry: ConfigEntry, + mock_added_config_entry: MockSchlageConfigEntry, mock_schlage: Mock, mock_lock: Mock, freezer: FrozenDateTimeFactory, @@ -38,7 +39,9 @@ async def test_lock_attributes( async def test_lock_services( - hass: HomeAssistant, mock_lock: Mock, mock_added_config_entry: ConfigEntry + hass: HomeAssistant, + mock_lock: Mock, + mock_added_config_entry: MockSchlageConfigEntry, ) -> None: """Test lock services.""" await hass.services.async_call( @@ -65,7 +68,7 @@ async def test_lock_services( async def test_changed_by( hass: HomeAssistant, mock_lock: Mock, - mock_added_config_entry: ConfigEntry, + mock_added_config_entry: MockSchlageConfigEntry, freezer: FrozenDateTimeFactory, ) -> None: """Test population of the changed_by attribute.""" diff --git a/tests/components/schlage/test_select.py b/tests/components/schlage/test_select.py index c27fd4c8813..59ff065d449 100644 --- a/tests/components/schlage/test_select.py +++ b/tests/components/schlage/test_select.py @@ -7,13 +7,16 @@ from homeassistant.components.select import ( DOMAIN as SELECT_DOMAIN, SERVICE_SELECT_OPTION, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ATTR_ENTITY_ID from homeassistant.core import HomeAssistant +from . import MockSchlageConfigEntry + async def test_select( - hass: HomeAssistant, mock_lock: Mock, mock_added_config_entry: ConfigEntry + hass: HomeAssistant, + mock_lock: Mock, + mock_added_config_entry: MockSchlageConfigEntry, ) -> None: """Test the auto-lock time select entity.""" entity_id = "select.vault_door_auto_lock_time" diff --git a/tests/components/schlage/test_sensor.py b/tests/components/schlage/test_sensor.py index 9fa90edecbb..9a489f6ff73 100644 --- a/tests/components/schlage/test_sensor.py +++ b/tests/components/schlage/test_sensor.py @@ -1,13 +1,14 @@ """Test schlage sensor.""" from homeassistant.components.sensor import SensorDeviceClass -from homeassistant.config_entries import ConfigEntry from homeassistant.const import PERCENTAGE from homeassistant.core import HomeAssistant +from . import MockSchlageConfigEntry + async def test_battery_sensor( - hass: HomeAssistant, mock_added_config_entry: ConfigEntry + hass: HomeAssistant, mock_added_config_entry: MockSchlageConfigEntry ) -> None: """Test the battery sensor.""" battery_sensor = hass.states.get("sensor.vault_door_battery") diff --git a/tests/components/schlage/test_switch.py b/tests/components/schlage/test_switch.py index 52b8da81670..fc5acc4399f 100644 --- a/tests/components/schlage/test_switch.py +++ b/tests/components/schlage/test_switch.py @@ -3,13 +3,16 @@ from unittest.mock import Mock 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 . import MockSchlageConfigEntry + async def test_beeper_services( - hass: HomeAssistant, mock_lock: Mock, mock_added_config_entry: ConfigEntry + hass: HomeAssistant, + mock_lock: Mock, + mock_added_config_entry: MockSchlageConfigEntry, ) -> None: """Test BeeperSwitch services.""" await hass.services.async_call( @@ -35,7 +38,9 @@ async def test_beeper_services( async def test_lock_and_leave_services( - hass: HomeAssistant, mock_lock: Mock, mock_added_config_entry: ConfigEntry + hass: HomeAssistant, + mock_lock: Mock, + mock_added_config_entry: MockSchlageConfigEntry, ) -> None: """Test LockAndLeaveSwitch services.""" await hass.services.async_call(