Store Schlage runtime data in entry.runtime_data (#131731)

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)

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(

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()]

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(

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(

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(

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(

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]

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 = []

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."""

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:

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."""

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:

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."""

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"

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")

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(