Use ConfigEntry.runtime_data in Nest (#131871)

homeassistant/components/nest/__init__.py

@@ -59,9 +59,7 @@ from .const import (
     CONF_SUBSCRIBER_ID,
     CONF_SUBSCRIBER_ID_IMPORTED,
     CONF_SUBSCRIPTION_NAME,
-    DATA_DEVICE_MANAGER,
     DATA_SDM,
-    DATA_SUBSCRIBER,
     DOMAIN,
 )
 from .events import EVENT_NAME_MAP, NEST_EVENT
@@ -72,6 +70,7 @@ from .media_source import (
     async_get_media_source_devices,
     async_get_transcoder,
 )
+from .types import NestConfigEntry, NestData
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -113,11 +112,8 @@ THUMBNAIL_SIZE_PX = 175
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up Nest components with dispatch between old/new flows."""
-    hass.data[DOMAIN] = {}
-
     hass.http.register_view(NestEventMediaView(hass))
     hass.http.register_view(NestEventMediaThumbnailView(hass))
-
     return True
 
 
@@ -128,12 +124,12 @@ class SignalUpdateCallback:
         self,
         hass: HomeAssistant,
         config_reload_cb: Callable[[], Awaitable[None]],
-        config_entry_id: str,
+        config_entry: NestConfigEntry,
     ) -> None:
         """Initialize EventCallback."""
         self._hass = hass
         self._config_reload_cb = config_reload_cb
-        self._config_entry_id = config_entry_id
+        self._config_entry = config_entry
 
     async def async_handle_event(self, event_message: EventMessage) -> None:
         """Process an incoming EventMessage."""
@@ -181,17 +177,17 @@ class SignalUpdateCallback:
                 message["zones"] = image_event.zones
             self._hass.bus.async_fire(NEST_EVENT, message)
 
-    def _supported_traits(self, device_id: str) -> list[TraitType]:
-        if not (
-            device_manager := self._hass.data[DOMAIN]
-            .get(self._config_entry_id, {})
-            .get(DATA_DEVICE_MANAGER)
-        ) or not (device := device_manager.devices.get(device_id)):
+    def _supported_traits(self, device_id: str) -> list[str]:
+        if (
+            not self._config_entry.runtime_data
+            or not (device_manager := self._config_entry.runtime_data.device_manager)
+            or not (device := device_manager.devices.get(device_id))
+        ):
             return []
         return list(device.traits)
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: NestConfigEntry) -> bool:
     """Set up Nest from a config entry with dispatch between old/new flows."""
     if DATA_SDM not in entry.data:
         hass.async_create_task(hass.config_entries.async_remove(entry.entry_id))
@@ -215,7 +211,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     async def async_config_reload() -> None:
         await hass.config_entries.async_reload(entry.entry_id)
 
-    update_callback = SignalUpdateCallback(hass, async_config_reload, entry.entry_id)
+    update_callback = SignalUpdateCallback(hass, async_config_reload, entry)
     subscriber.set_update_callback(update_callback.async_handle_event)
     try:
         await subscriber.start_async()
@@ -245,11 +241,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     entry.async_on_unload(
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, on_hass_stop)
     )
-
-    hass.data[DOMAIN][entry.entry_id] = {
-        DATA_SUBSCRIBER: subscriber,
-        DATA_DEVICE_MANAGER: device_manager,
-    }
+    entry.runtime_data = NestData(
+        subscriber=subscriber,
+        device_manager=device_manager,
+    )
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -262,13 +257,9 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         # Legacy API
         return True
     _LOGGER.debug("Stopping nest subscriber")
-    subscriber = hass.data[DOMAIN][entry.entry_id][DATA_SUBSCRIBER]
+    subscriber = entry.runtime_data.subscriber
     subscriber.stop_async()
-    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)
 
 
 async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:

homeassistant/components/nest/camera.py

@@ -17,7 +17,6 @@ from google_nest_sdm.camera_traits import (
     WebRtcStream,
 )
 from google_nest_sdm.device import Device
-from google_nest_sdm.device_manager import DeviceManager
 from google_nest_sdm.exceptions import ApiException
 from webrtc_models import RTCIceCandidateInit
 
@@ -29,15 +28,14 @@ from homeassistant.components.camera import (
     WebRTCSendMessage,
 )
 from homeassistant.components.stream import CONF_EXTRA_PART_WAIT_TIME
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_track_point_in_utc_time
 from homeassistant.util.dt import utcnow
 
-from .const import DATA_DEVICE_MANAGER, DOMAIN
 from .device_info import NestDeviceInfo
+from .types import NestConfigEntry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -53,15 +51,12 @@ BACKOFF_MULTIPLIER = 1.5
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant, entry: NestConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the cameras."""
 
-    device_manager: DeviceManager = hass.data[DOMAIN][entry.entry_id][
-        DATA_DEVICE_MANAGER
-    ]
     entities: list[NestCameraBaseEntity] = []
-    for device in device_manager.devices.values():
+    for device in entry.runtime_data.device_manager.devices.values():
         if (live_stream := device.traits.get(CameraLiveStreamTrait.NAME)) is None:
             continue
         if StreamingProtocol.WEB_RTC in live_stream.supported_protocols:

homeassistant/components/nest/climate.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 from typing import Any, cast
 
 from google_nest_sdm.device import Device
-from google_nest_sdm.device_manager import DeviceManager
 from google_nest_sdm.device_traits import FanTrait, TemperatureTrait
 from google_nest_sdm.exceptions import ApiException
 from google_nest_sdm.thermostat_traits import (
@@ -28,14 +27,13 @@ from homeassistant.components.climate import (
     HVACAction,
     HVACMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DATA_DEVICE_MANAGER, DOMAIN
 from .device_info import NestDeviceInfo
+from .types import NestConfigEntry
 
 # Mapping for sdm.devices.traits.ThermostatMode mode field
 THERMOSTAT_MODE_MAP: dict[str, HVACMode] = {
@@ -78,17 +76,13 @@ MIN_TEMP_RANGE = 1.66667
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant, entry: NestConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the client entities."""
 
-    device_manager: DeviceManager = hass.data[DOMAIN][entry.entry_id][
-        DATA_DEVICE_MANAGER
-    ]
-
     async_add_entities(
         ThermostatEntity(device)
-        for device in device_manager.devices.values()
+        for device in entry.runtime_data.device_manager.devices.values()
         if ThermostatHvacTrait.NAME in device.traits
     )
 

homeassistant/components/nest/const.py

@@ -2,8 +2,6 @@
 
 DOMAIN = "nest"
 DATA_SDM = "sdm"
-DATA_SUBSCRIBER = "subscriber"
-DATA_DEVICE_MANAGER = "device_manager"
 
 WEB_AUTH_DOMAIN = DOMAIN
 INSTALLED_AUTH_DOMAIN = f"{DOMAIN}.installed"

homeassistant/components/nest/device_info.py

@@ -7,11 +7,12 @@ from collections.abc import Mapping
 from google_nest_sdm.device import Device
 from google_nest_sdm.device_traits import ConnectivityTrait, InfoTrait
 
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo
 
-from .const import CONNECTIVITY_TRAIT_OFFLINE, DATA_DEVICE_MANAGER, DOMAIN
+from .const import CONNECTIVITY_TRAIT_OFFLINE, DOMAIN
 
 DEVICE_TYPE_MAP: dict[str, str] = {
     "sdm.devices.types.CAMERA": "Camera",
@@ -81,14 +82,12 @@ class NestDeviceInfo:
 @callback
 def async_nest_devices(hass: HomeAssistant) -> Mapping[str, Device]:
     """Return a mapping of all nest devices for all config entries."""
-    devices = {}
-    for entry_id in hass.data[DOMAIN]:
-        if not (device_manager := hass.data[DOMAIN][entry_id].get(DATA_DEVICE_MANAGER)):
-            continue
-        devices.update(
-            {device.name: device for device in device_manager.devices.values()}
-        )
-    return devices
+    return {
+        device.name: device
+        for config_entry in hass.config_entries.async_entries(DOMAIN)
+        if config_entry.state == ConfigEntryState.LOADED
+        for device in config_entry.runtime_data.device_manager.devices.values()
+    }
 
 
 @callback

homeassistant/components/nest/diagnostics.py

@@ -5,46 +5,26 @@ from __future__ import annotations
 from typing import Any
 
 from google_nest_sdm import diagnostics
-from google_nest_sdm.device import Device
-from google_nest_sdm.device_manager import DeviceManager
 from google_nest_sdm.device_traits import InfoTrait
 
 from homeassistant.components.camera import diagnostics as camera_diagnostics
-from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceEntry
 
-from .const import DATA_DEVICE_MANAGER, DATA_SDM, DOMAIN
+from .types import NestConfigEntry
 
 REDACT_DEVICE_TRAITS = {InfoTrait.NAME}
 
 
-@callback
-def _async_get_nest_devices(
-    hass: HomeAssistant, config_entry: ConfigEntry
-) -> dict[str, Device]:
-    """Return dict of available devices."""
-    if DATA_SDM not in config_entry.data:
-        return {}
-
-    if (
-        config_entry.entry_id not in hass.data[DOMAIN]
-        or DATA_DEVICE_MANAGER not in hass.data[DOMAIN][config_entry.entry_id]
-    ):
-        return {}
-
-    device_manager: DeviceManager = hass.data[DOMAIN][config_entry.entry_id][
-        DATA_DEVICE_MANAGER
-    ]
-    return device_manager.devices
-
-
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: NestConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    nest_devices = _async_get_nest_devices(hass, config_entry)
-    if not nest_devices:
+    if (
+        not hasattr(config_entry, "runtime_data")
+        or not config_entry.runtime_data
+        or not (nest_devices := config_entry.runtime_data.device_manager.devices)
+    ):
         return {}
     data: dict[str, Any] = {
         **diagnostics.get_diagnostics(),
@@ -62,11 +42,11 @@ async def async_get_config_entry_diagnostics(
 
 async def async_get_device_diagnostics(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: NestConfigEntry,
     device: DeviceEntry,
 ) -> dict[str, Any]:
     """Return diagnostics for a device."""
-    nest_devices = _async_get_nest_devices(hass, config_entry)
+    nest_devices = config_entry.runtime_data.device_manager.devices
     nest_device_id = next(iter(device.identifiers))[1]
     nest_device = nest_devices.get(nest_device_id)
     return nest_device.get_diagnostics() if nest_device else {}

homeassistant/components/nest/event.py

@@ -4,7 +4,6 @@ from dataclasses import dataclass
 import logging
 
 from google_nest_sdm.device import Device
-from google_nest_sdm.device_manager import DeviceManager
 from google_nest_sdm.event import EventMessage, EventType
 from google_nest_sdm.traits import TraitType
 
@@ -13,11 +12,9 @@ from homeassistant.components.event import (
     EventEntity,
     EventEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DATA_DEVICE_MANAGER, DOMAIN
 from .device_info import NestDeviceInfo
 from .events import (
     EVENT_CAMERA_MOTION,
@@ -26,6 +23,7 @@ from .events import (
     EVENT_DOORBELL_CHIME,
     EVENT_NAME_MAP,
 )
+from .types import NestConfigEntry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -68,16 +66,12 @@ ENTITY_DESCRIPTIONS = [
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant, entry: NestConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the sensors."""
-
-    device_manager: DeviceManager = hass.data[DOMAIN][entry.entry_id][
-        DATA_DEVICE_MANAGER
-    ]
     async_add_entities(
         NestTraitEventEntity(desc, device)
-        for device in device_manager.devices.values()
+        for device in entry.runtime_data.device_manager.devices.values()
         for desc in ENTITY_DESCRIPTIONS
         if any(trait in device.traits for trait in desc.trait_types)
     )

homeassistant/components/nest/sensor.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 import logging
 
 from google_nest_sdm.device import Device
-from google_nest_sdm.device_manager import DeviceManager
 from google_nest_sdm.device_traits import HumidityTrait, TemperatureTrait
 
 from homeassistant.components.sensor import (
@@ -13,13 +12,12 @@ from homeassistant.components.sensor import (
     SensorEntity,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE, UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DATA_DEVICE_MANAGER, DOMAIN
 from .device_info import NestDeviceInfo
+from .types import NestConfigEntry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -33,15 +31,12 @@ DEVICE_TYPE_MAP = {
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant, entry: NestConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the sensors."""
 
-    device_manager: DeviceManager = hass.data[DOMAIN][entry.entry_id][
-        DATA_DEVICE_MANAGER
-    ]
     entities: list[SensorEntity] = []
-    for device in device_manager.devices.values():
+    for device in entry.runtime_data.device_manager.devices.values():
         if TemperatureTrait.NAME in device.traits:
             entities.append(TemperatureSensor(device))
         if HumidityTrait.NAME in device.traits:

homeassistant/components/nest/types.py

@@ -0,0 +1,19 @@
+"""Type definitions for Nest."""
+
+from dataclasses import dataclass
+
+from google_nest_sdm.device_manager import DeviceManager
+from google_nest_sdm.google_nest_subscriber import GoogleNestSubscriber
+
+from homeassistant.config_entries import ConfigEntry
+
+
+@dataclass
+class NestData:
+    """Data for the Nest integration."""
+
+    subscriber: GoogleNestSubscriber
+    device_manager: DeviceManager
+
+
+type NestConfigEntry = ConfigEntry[NestData]