Use runtime_data in guardian (#144344)

* Use runtime_data in guardian

* Adjust tests

homeassistant/components/guardian/__init__.py

@@ -40,12 +40,14 @@ PLATFORMS = [
     Platform.VALVE,
 ]
 
+type GuardianConfigEntry = ConfigEntry[GuardianData]
+
 
 @dataclass
 class GuardianData:
-    """Define an object to be stored in `hass.data`."""
+    """Define an object to be stored in `entry.runtime_data`."""
 
-    entry: ConfigEntry
+    entry: GuardianConfigEntry
     client: Client
     valve_controller_coordinators: dict[str, GuardianDataUpdateCoordinator]
     paired_sensor_manager: PairedSensorManager
@@ -57,7 +59,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: GuardianConfigEntry) -> bool:
     """Set up Elexa Guardian from a config entry."""
     client = Client(entry.data[CONF_IP_ADDRESS], port=entry.data[CONF_PORT])
 
@@ -108,8 +110,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
     await paired_sensor_manager.async_initialize()
 
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = GuardianData(
+    entry.runtime_data = GuardianData(
         entry=entry,
         client=client,
         valve_controller_coordinators=valve_controller_coordinators,
@@ -122,13 +123,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: GuardianConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
 class PairedSensorManager:
@@ -137,7 +134,7 @@ class PairedSensorManager:
     def __init__(
         self,
         hass: HomeAssistant,
-        entry: ConfigEntry,
+        entry: GuardianConfigEntry,
         client: Client,
         api_lock: asyncio.Lock,
         sensor_pair_dump_coordinator: GuardianDataUpdateCoordinator,

homeassistant/components/guardian/binary_sensor.py

@@ -12,17 +12,15 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from . import GuardianData
+from . import GuardianConfigEntry
 from .const import (
     API_SYSTEM_ONBOARD_SENSOR_STATUS,
     CONF_UID,
-    DOMAIN,
     SIGNAL_PAIRED_SENSOR_COORDINATOR_ADDED,
 )
 from .coordinator import GuardianDataUpdateCoordinator
@@ -87,11 +85,11 @@ VALVE_CONTROLLER_DESCRIPTIONS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: GuardianConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Guardian switches based on a config entry."""
-    data: GuardianData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
     uid = entry.data[CONF_UID]
 
     async_finish_entity_domain_replacements(
@@ -151,7 +149,7 @@ class PairedSensorBinarySensor(PairedSensorEntity, BinarySensorEntity):
 
     def __init__(
         self,
-        entry: ConfigEntry,
+        entry: GuardianConfigEntry,
         coordinator: GuardianDataUpdateCoordinator,
         description: BinarySensorEntityDescription,
     ) -> None:
@@ -173,7 +171,7 @@ class ValveControllerBinarySensor(ValveControllerEntity, BinarySensorEntity):
 
     def __init__(
         self,
-        entry: ConfigEntry,
+        entry: GuardianConfigEntry,
         coordinators: dict[str, GuardianDataUpdateCoordinator],
         description: ValveControllerBinarySensorDescription,
     ) -> None:

homeassistant/components/guardian/button.py

@@ -12,14 +12,13 @@ from homeassistant.components.button import (
     ButtonEntity,
     ButtonEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from . import GuardianData
-from .const import API_SYSTEM_DIAGNOSTICS, DOMAIN
+from . import GuardianConfigEntry, GuardianData
+from .const import API_SYSTEM_DIAGNOSTICS
 from .entity import ValveControllerEntity, ValveControllerEntityDescription
 from .util import convert_exceptions_to_homeassistant_error
 
@@ -69,11 +68,11 @@ BUTTON_DESCRIPTIONS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: GuardianConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Guardian buttons based on a config entry."""
-    data: GuardianData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     async_add_entities(
         GuardianButton(entry, data, description) for description in BUTTON_DESCRIPTIONS
@@ -90,7 +89,7 @@ class GuardianButton(ValveControllerEntity, ButtonEntity):
 
     def __init__(
         self,
-        entry: ConfigEntry,
+        entry: GuardianConfigEntry,
         data: GuardianData,
         description: ValveControllerButtonDescription,
     ) -> None:

homeassistant/components/guardian/coordinator.py

@@ -5,18 +5,20 @@ from __future__ import annotations
 import asyncio
 from collections.abc import Callable, Coroutine
 from datetime import timedelta
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from aioguardian import Client
 from aioguardian.errors import GuardianError
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import LOGGER
 
+if TYPE_CHECKING:
+    from . import GuardianConfigEntry
+
 DEFAULT_UPDATE_INTERVAL = timedelta(seconds=30)
 
 SIGNAL_REBOOT_REQUESTED = "guardian_reboot_requested_{0}"
@@ -25,13 +27,13 @@ SIGNAL_REBOOT_REQUESTED = "guardian_reboot_requested_{0}"
 class GuardianDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]):
     """Define an extended DataUpdateCoordinator with some Guardian goodies."""
 
-    config_entry: ConfigEntry
+    config_entry: GuardianConfigEntry
 
     def __init__(
         self,
         hass: HomeAssistant,
         *,
-        entry: ConfigEntry,
+        entry: GuardianConfigEntry,
         client: Client,
         api_name: str,
         api_coro: Callable[[], Coroutine[Any, Any, dict[str, Any]]],

homeassistant/components/guardian/diagnostics.py

@@ -5,12 +5,11 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_UNIQUE_ID
 from homeassistant.core import HomeAssistant
 
-from . import GuardianData
-from .const import CONF_UID, DOMAIN
+from . import GuardianConfigEntry
+from .const import CONF_UID
 
 CONF_BSSID = "bssid"
 CONF_PAIRED_UIDS = "paired_uids"
@@ -29,10 +28,10 @@ TO_REDACT = {
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: GuardianConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    data: GuardianData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     return {
         "entry": async_redact_data(entry.as_dict(), TO_REDACT),

homeassistant/components/guardian/entity.py

@@ -4,11 +4,11 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import EntityDescription
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from . import GuardianConfigEntry
 from .const import API_SYSTEM_DIAGNOSTICS, CONF_UID, DOMAIN
 from .coordinator import GuardianDataUpdateCoordinator
 
@@ -32,7 +32,7 @@ class PairedSensorEntity(GuardianEntity):
 
     def __init__(
         self,
-        entry: ConfigEntry,
+        entry: GuardianConfigEntry,
         coordinator: GuardianDataUpdateCoordinator,
         description: EntityDescription,
     ) -> None:
@@ -62,7 +62,7 @@ class ValveControllerEntity(GuardianEntity):
 
     def __init__(
         self,
-        entry: ConfigEntry,
+        entry: GuardianConfigEntry,
         coordinators: dict[str, GuardianDataUpdateCoordinator],
         description: ValveControllerEntityDescription,
     ) -> None:

homeassistant/components/guardian/sensor.py

@@ -12,7 +12,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     EntityCategory,
     UnitOfElectricCurrent,
@@ -25,13 +24,12 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
-from . import GuardianData
+from . import GuardianConfigEntry
 from .const import (
     API_SYSTEM_DIAGNOSTICS,
     API_SYSTEM_ONBOARD_SENSOR_STATUS,
     API_VALVE_STATUS,
     CONF_UID,
-    DOMAIN,
     SIGNAL_PAIRED_SENSOR_COORDINATOR_ADDED,
 )
 from .entity import (
@@ -138,11 +136,11 @@ VALVE_CONTROLLER_DESCRIPTIONS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: GuardianConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Guardian switches based on a config entry."""
-    data: GuardianData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     @callback
     def add_new_paired_sensor(uid: str) -> None:

homeassistant/components/guardian/services.py

@@ -22,7 +22,7 @@ from homeassistant.helpers import config_validation as cv, device_registry as dr
 from .const import CONF_UID, DOMAIN
 
 if TYPE_CHECKING:
-    from . import GuardianData
+    from . import GuardianConfigEntry, GuardianData
 
 SERVICE_NAME_PAIR_SENSOR = "pair_sensor"
 SERVICE_NAME_UNPAIR_SENSOR = "unpair_sensor"
@@ -58,7 +58,7 @@ SERVICE_UPGRADE_FIRMWARE_SCHEMA = vol.Schema(
 
 
 @callback
-def async_get_entry_id_for_service_call(call: ServiceCall) -> str:
+def async_get_entry_id_for_service_call(call: ServiceCall) -> GuardianConfigEntry:
     """Get the entry ID related to a service call (by device ID)."""
     device_id = call.data[CONF_DEVICE_ID]
     device_registry = dr.async_get(call.hass)
@@ -70,7 +70,7 @@ def async_get_entry_id_for_service_call(call: ServiceCall) -> str:
         if (entry := call.hass.config_entries.async_get_entry(entry_id)) is None:
             continue
         if entry.domain == DOMAIN:
-            return entry_id
+            return entry
 
     raise ValueError(f"No config entry for device ID: {device_id}")
 
@@ -83,8 +83,7 @@ def call_with_data(
 
     async def wrapper(call: ServiceCall) -> None:
         """Wrap the service function."""
-        entry_id = async_get_entry_id_for_service_call(call)
-        data = call.hass.data[DOMAIN][entry_id]
+        data = async_get_entry_id_for_service_call(call).runtime_data
 
         try:
             async with data.client:

homeassistant/components/guardian/switch.py

@@ -9,13 +9,12 @@ from typing import Any
 from aioguardian import Client
 
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from . import GuardianData
-from .const import API_VALVE_STATUS, API_WIFI_STATUS, DOMAIN
+from . import GuardianConfigEntry, GuardianData
+from .const import API_VALVE_STATUS, API_WIFI_STATUS
 from .entity import ValveControllerEntity, ValveControllerEntityDescription
 from .util import convert_exceptions_to_homeassistant_error
 from .valve import GuardianValveState
@@ -111,11 +110,11 @@ VALVE_CONTROLLER_DESCRIPTIONS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: GuardianConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Guardian switches based on a config entry."""
-    data: GuardianData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     async_add_entities(
         ValveControllerSwitch(entry, data, description)
@@ -130,7 +129,7 @@ class ValveControllerSwitch(ValveControllerEntity, SwitchEntity):
 
     def __init__(
         self,
-        entry: ConfigEntry,
+        entry: GuardianConfigEntry,
         data: GuardianData,
         description: ValveControllerSwitchDescription,
     ) -> None:

homeassistant/components/guardian/util.py

@@ -10,7 +10,6 @@ from typing import TYPE_CHECKING, Any, Concatenate
 
 from aioguardian.errors import GuardianError
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
@@ -18,6 +17,7 @@ from homeassistant.helpers import entity_registry as er
 from .const import LOGGER
 
 if TYPE_CHECKING:
+    from . import GuardianConfigEntry
     from .entity import GuardianEntity
 
 DEFAULT_UPDATE_INTERVAL = timedelta(seconds=30)
@@ -36,7 +36,7 @@ class EntityDomainReplacementStrategy:
 @callback
 def async_finish_entity_domain_replacements(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: GuardianConfigEntry,
     entity_replacement_strategies: Iterable[EntityDomainReplacementStrategy],
 ) -> None:
     """Remove old entities and create a repairs issue with info on their replacement."""

homeassistant/components/guardian/valve.py

@@ -15,12 +15,11 @@ from homeassistant.components.valve import (
     ValveEntityDescription,
     ValveEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from . import GuardianData
-from .const import API_VALVE_STATUS, DOMAIN
+from . import GuardianConfigEntry, GuardianData
+from .const import API_VALVE_STATUS
 from .entity import ValveControllerEntity, ValveControllerEntityDescription
 from .util import convert_exceptions_to_homeassistant_error
 
@@ -110,11 +109,11 @@ VALVE_CONTROLLER_DESCRIPTIONS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: GuardianConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Guardian switches based on a config entry."""
-    data: GuardianData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     async_add_entities(
         ValveControllerValve(entry, data, description)
@@ -132,7 +131,7 @@ class ValveControllerValve(ValveControllerEntity, ValveEntity):
 
     def __init__(
         self,
-        entry: ConfigEntry,
+        entry: GuardianConfigEntry,
         data: GuardianData,
         description: ValveControllerValveDescription,
     ) -> None:

tests/components/guardian/test_diagnostics.py

@@ -1,7 +1,7 @@
 """Test Guardian diagnostics."""
 
 from homeassistant.components.diagnostics import REDACTED
-from homeassistant.components.guardian import DOMAIN, GuardianData
+from homeassistant.components.guardian import GuardianData
 from homeassistant.core import HomeAssistant
 
 from tests.common import ANY, MockConfigEntry
@@ -16,7 +16,7 @@ async def test_entry_diagnostics(
     setup_guardian: None,  # relies on config_entry fixture
 ) -> None:
     """Test config entry diagnostics."""
-    data: GuardianData = hass.data[DOMAIN][config_entry.entry_id]
+    data: GuardianData = config_entry.runtime_data
 
     # Simulate the pairing of a paired sensor:
     await data.paired_sensor_manager.async_pair_sensor("AABBCCDDEEFF")