Store runtime data inside the config entry in Sense (#119740)

homeassistant/components/sense/__init__.py

@@ -1,7 +1,9 @@
 """Support for monitoring a Sense energy sensor."""
 
+from dataclasses import dataclass
 from datetime import timedelta
 import logging
+from typing import Any
 
 from sense_energy import (
     ASyncSenseable,
@@ -25,20 +27,16 @@ from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, Upda
 
 from .const import (
     ACTIVE_UPDATE_RATE,
-    DOMAIN,
     SENSE_CONNECT_EXCEPTIONS,
-    SENSE_DATA,
     SENSE_DEVICE_UPDATE,
-    SENSE_DEVICES_DATA,
-    SENSE_DISCOVERED_DEVICES_DATA,
     SENSE_TIMEOUT_EXCEPTIONS,
-    SENSE_TRENDS_COORDINATOR,
     SENSE_WEBSOCKET_EXCEPTIONS,
 )
 
 _LOGGER = logging.getLogger(__name__)
 
 PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR]
+type SenseConfigEntry = ConfigEntry[SenseData]
 
 
 class SenseDevicesData:
@@ -57,7 +55,17 @@ class SenseDevicesData:
         return self._data_by_device.get(sense_device_id)
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+@dataclass(kw_only=True, slots=True)
+class SenseData:
+    """Sense data type."""
+
+    data: ASyncSenseable
+    device_data: SenseDevicesData
+    trends: DataUpdateCoordinator[None]
+    discovered: list[dict[str, Any]]
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: SenseConfigEntry) -> bool:
     """Set up Sense from a config entry."""
 
     entry_data = entry.data
@@ -91,7 +99,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     except SENSE_CONNECT_EXCEPTIONS as err:
         raise ConfigEntryNotReady(str(err)) from err
 
-    sense_devices_data = SenseDevicesData()
     try:
         sense_discovered_devices = await gateway.get_discovered_device_data()
         await gateway.update_realtime()
@@ -132,12 +139,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         "sense.trends-coordinator-refresh",
     )
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {
-        SENSE_DATA: gateway,
-        SENSE_DEVICES_DATA: sense_devices_data,
-        SENSE_TRENDS_COORDINATOR: trends_coordinator,
-        SENSE_DISCOVERED_DEVICES_DATA: sense_discovered_devices,
-    }
+    entry.runtime_data = SenseData(
+        data=gateway,
+        device_data=SenseDevicesData(),
+        trends=trends_coordinator,
+        discovered=sense_discovered_devices,
+    )
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -152,7 +159,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
         data = gateway.get_realtime()
         if "devices" in data:
-            sense_devices_data.set_devices_data(data["devices"])
+            entry.runtime_data.device_data.set_devices_data(data["devices"])
         async_dispatcher_send(hass, f"{SENSE_DEVICE_UPDATE}-{gateway.sense_monitor_id}")
 
     remove_update_callback = async_track_time_interval(
@@ -173,9 +180,6 @@ 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: SenseConfigEntry) -> 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)

homeassistant/components/sense/binary_sensor.py

@@ -6,40 +6,29 @@ from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import (
-    ATTRIBUTION,
-    DOMAIN,
-    MDI_ICONS,
-    SENSE_DATA,
-    SENSE_DEVICE_UPDATE,
-    SENSE_DEVICES_DATA,
-    SENSE_DISCOVERED_DEVICES_DATA,
-)
+from . import SenseConfigEntry
+from .const import ATTRIBUTION, DOMAIN, MDI_ICONS, SENSE_DEVICE_UPDATE
 
 _LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: SenseConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Sense binary sensor."""
-    data = hass.data[DOMAIN][config_entry.entry_id][SENSE_DATA]
-    sense_devices_data = hass.data[DOMAIN][config_entry.entry_id][SENSE_DEVICES_DATA]
-    sense_monitor_id = data.sense_monitor_id
+    sense_monitor_id = config_entry.runtime_data.data.sense_monitor_id
 
-    sense_devices = hass.data[DOMAIN][config_entry.entry_id][
-        SENSE_DISCOVERED_DEVICES_DATA
-    ]
+    sense_devices = config_entry.runtime_data.discovered
+    device_data = config_entry.runtime_data.device_data
     devices = [
-        SenseDevice(sense_devices_data, device, sense_monitor_id)
+        SenseDevice(device_data, device, sense_monitor_id)
         for device in sense_devices
         if device["tags"]["DeviceListAllowed"] == "true"
     ]

homeassistant/components/sense/const.py

@@ -12,11 +12,7 @@ DOMAIN = "sense"
 DEFAULT_TIMEOUT = 30
 ACTIVE_UPDATE_RATE = 60
 DEFAULT_NAME = "Sense"
-SENSE_DATA = "sense_data"
 SENSE_DEVICE_UPDATE = "sense_devices_update"
-SENSE_DEVICES_DATA = "sense_devices_data"
-SENSE_DISCOVERED_DEVICES_DATA = "sense_discovered_devices"
-SENSE_TRENDS_COORDINATOR = "sense_trends_coordinator"
 
 ACTIVE_NAME = "Energy"
 ACTIVE_TYPE = "active"

homeassistant/components/sense/sensor.py

@@ -5,7 +5,6 @@ from homeassistant.components.sensor import (
     SensorEntity,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     PERCENTAGE,
     UnitOfElectricPotential,
@@ -18,6 +17,7 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from . import SenseConfigEntry
 from .const import (
     ACTIVE_NAME,
     ACTIVE_TYPE,
@@ -34,11 +34,7 @@ from .const import (
     PRODUCTION_NAME,
     PRODUCTION_PCT_ID,
     PRODUCTION_PCT_NAME,
-    SENSE_DATA,
     SENSE_DEVICE_UPDATE,
-    SENSE_DEVICES_DATA,
-    SENSE_DISCOVERED_DEVICES_DATA,
-    SENSE_TRENDS_COORDINATOR,
     SOLAR_POWERED_ID,
     SOLAR_POWERED_NAME,
     TO_GRID_ID,
@@ -87,26 +83,23 @@ def sense_to_mdi(sense_icon):
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: SenseConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Sense sensor."""
-    base_data = hass.data[DOMAIN][config_entry.entry_id]
-    data = base_data[SENSE_DATA]
-    sense_devices_data = base_data[SENSE_DEVICES_DATA]
-    trends_coordinator = base_data[SENSE_TRENDS_COORDINATOR]
+    data = config_entry.runtime_data.data
+    trends_coordinator = config_entry.runtime_data.trends
 
     # Request only in case it takes longer
     # than 60s
     await trends_coordinator.async_request_refresh()
 
     sense_monitor_id = data.sense_monitor_id
-    sense_devices = hass.data[DOMAIN][config_entry.entry_id][
-        SENSE_DISCOVERED_DEVICES_DATA
-    ]
+    sense_devices = config_entry.runtime_data.discovered
+    device_data = config_entry.runtime_data.device_data
 
     entities: list[SensorEntity] = [
-        SenseEnergyDevice(sense_devices_data, device, sense_monitor_id)
+        SenseEnergyDevice(device_data, device, sense_monitor_id)
         for device in sense_devices
         if device["tags"]["DeviceListAllowed"] == "true"
     ]