Use runtime_data in flume (#137660)

homeassistant/components/flume/__init__.py

@@ -7,7 +7,7 @@ from requests import Session
 from requests.exceptions import RequestException
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.const import (
     CONF_CLIENT_ID,
     CONF_CLIENT_SECRET,
@@ -24,16 +24,12 @@ from homeassistant.core import (
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers.selector import ConfigEntrySelector
 
-from .const import (
+from .const import BASE_TOKEN_FILENAME, DOMAIN, PLATFORMS
-    BASE_TOKEN_FILENAME,
+from .coordinator import (
-    DOMAIN,
+    FlumeConfigEntry,
-    FLUME_AUTH,
+    FlumeNotificationDataUpdateCoordinator,
-    FLUME_DEVICES,
+    FlumeRuntimeData,
-    FLUME_HTTP_SESSION,
-    FLUME_NOTIFICATIONS_COORDINATOR,
-    PLATFORMS,
 )
-from .coordinator import FlumeNotificationDataUpdateCoordinator
 
 SERVICE_LIST_NOTIFICATIONS = "list_notifications"
 CONF_CONFIG_ENTRY = "config_entry"
@@ -76,7 +72,7 @@ def _setup_entry(
     return flume_auth, flume_devices, http_session
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: FlumeConfigEntry) -> bool:
     """Set up flume from a config entry."""
 
     flume_auth, flume_devices, http_session = await hass.async_add_executor_job(
@@ -86,12 +82,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass=hass, auth=flume_auth
     )
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {
+    entry.runtime_data = FlumeRuntimeData(
-        FLUME_DEVICES: flume_devices,
+        devices=flume_devices,
-        FLUME_AUTH: flume_auth,
+        auth=flume_auth,
-        FLUME_HTTP_SESSION: http_session,
+        http_session=http_session,
-        FLUME_NOTIFICATIONS_COORDINATOR: notification_coordinator,
+        notifications_coordinator=notification_coordinator,
-    }
+    )
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     setup_service(hass)
@@ -99,16 +95,10 @@ 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: FlumeConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    entry.runtime_data.http_session.close()
-
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    hass.data[DOMAIN][entry.entry_id][FLUME_HTTP_SESSION].close()
-
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
 
 
 def setup_service(hass: HomeAssistant) -> None:
@@ -118,15 +108,13 @@ def setup_service(hass: HomeAssistant) -> None:
     def list_notifications(call: ServiceCall) -> ServiceResponse:
         """Return the user notifications."""
         entry_id: str = call.data[CONF_CONFIG_ENTRY]
-        entry: ConfigEntry | None = hass.config_entries.async_get_entry(entry_id)
+        entry: FlumeConfigEntry | None = hass.config_entries.async_get_entry(entry_id)
         if not entry:
             raise ValueError(f"Invalid config entry: {entry_id}")
-        if not (flume_domain_data := hass.data[DOMAIN].get(entry_id)):
+        if not entry.state == ConfigEntryState.LOADED:
             raise ValueError(f"Config entry not loaded: {entry_id}")
         return {
-            "notifications": flume_domain_data[
+            "notifications": entry.runtime_data.notifications_coordinator.notifications  # type: ignore[dict-item]
-                FLUME_NOTIFICATIONS_COORDINATOR
-            ].notifications
         }
 
     hass.services.async_register(

homeassistant/components/flume/binary_sensor.py

@@ -9,15 +9,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,
-    FLUME_DEVICES,
-    FLUME_NOTIFICATIONS_COORDINATOR,
     FLUME_TYPE_BRIDGE,
     FLUME_TYPE_SENSOR,
     KEY_DEVICE_ID,
@@ -29,6 +25,7 @@ from .const import (
     NOTIFICATION_LOW_BATTERY,
 )
 from .coordinator import (
+    FlumeConfigEntry,
     FlumeDeviceConnectionUpdateCoordinator,
     FlumeNotificationDataUpdateCoordinator,
 )
@@ -71,12 +68,12 @@ FLUME_BINARY_NOTIFICATION_SENSORS: tuple[FlumeBinarySensorEntityDescription, ...
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: FlumeConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up a Flume binary sensor.."""
-    flume_domain_data = hass.data[DOMAIN][config_entry.entry_id]
+    flume_domain_data = config_entry.runtime_data
-    flume_devices = flume_domain_data[FLUME_DEVICES]
+    flume_devices = flume_domain_data.devices
 
     flume_entity_list: list[
         FlumeNotificationBinarySensor | FlumeConnectionBinarySensor
@@ -85,7 +82,7 @@ async def async_setup_entry(
     connection_coordinator = FlumeDeviceConnectionUpdateCoordinator(
         hass=hass, flume_devices=flume_devices
     )
-    notification_coordinator = flume_domain_data[FLUME_NOTIFICATIONS_COORDINATOR]
+    notification_coordinator = flume_domain_data.notifications_coordinator
     flume_devices = get_valid_flume_devices(flume_devices)
     for device in flume_devices:
         device_id = device[KEY_DEVICE_ID]

homeassistant/components/flume/const.py

@@ -26,12 +26,6 @@ _LOGGER = logging.getLogger(__package__)
 FLUME_TYPE_BRIDGE = 1
 FLUME_TYPE_SENSOR = 2
 
-
-FLUME_AUTH = "flume_auth"
-FLUME_HTTP_SESSION = "http_session"
-FLUME_DEVICES = "devices"
-FLUME_NOTIFICATIONS_COORDINATOR = "notifications_coordinator"
-
 CONF_TOKEN_FILE = "token_filename"
 BASE_TOKEN_FILENAME = "FLUME_TOKEN_FILE"
 

homeassistant/components/flume/coordinator.py

@@ -2,11 +2,14 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import Any
 
 import pyflume
 from pyflume import FlumeAuth, FlumeData, FlumeDeviceList
+from requests import Session
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
@@ -19,6 +22,19 @@ from .const import (
 )
 
 
+@dataclass
+class FlumeRuntimeData:
+    """Runtime data for the Flume config entry."""
+
+    devices: FlumeDeviceList
+    auth: FlumeAuth
+    http_session: Session
+    notifications_coordinator: FlumeNotificationDataUpdateCoordinator
+
+
+type FlumeConfigEntry = ConfigEntry[FlumeRuntimeData]
+
+
 class FlumeDeviceDataUpdateCoordinator(DataUpdateCoordinator[None]):
     """Data update coordinator for an individual flume device."""
 

homeassistant/components/flume/sensor.py

@@ -11,7 +11,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfVolume
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -19,10 +18,6 @@ from homeassistant.helpers.typing import StateType
 
 from .const import (
     DEVICE_SCAN_INTERVAL,
-    DOMAIN,
-    FLUME_AUTH,
-    FLUME_DEVICES,
-    FLUME_HTTP_SESSION,
     FLUME_TYPE_SENSOR,
     KEY_DEVICE_ID,
     KEY_DEVICE_LOCATION,
@@ -30,7 +25,7 @@ from .const import (
     KEY_DEVICE_LOCATION_TIMEZONE,
     KEY_DEVICE_TYPE,
 )
-from .coordinator import FlumeDeviceDataUpdateCoordinator
+from .coordinator import FlumeConfigEntry, FlumeDeviceDataUpdateCoordinator
 from .entity import FlumeEntity
 from .util import get_valid_flume_devices
 
@@ -112,15 +107,15 @@ def make_flume_datas(
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: FlumeConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Flume sensor."""
 
-    flume_domain_data = hass.data[DOMAIN][config_entry.entry_id]
+    flume_domain_data = config_entry.runtime_data
-    flume_devices = flume_domain_data[FLUME_DEVICES]
+    flume_devices = flume_domain_data.devices
-    flume_auth: FlumeAuth = flume_domain_data[FLUME_AUTH]
+    flume_auth = flume_domain_data.auth
-    http_session: Session = flume_domain_data[FLUME_HTTP_SESSION]
+    http_session = flume_domain_data.http_session
     flume_devices = [
         device
         for device in get_valid_flume_devices(flume_devices)