Use ConfigEntry runtime_data in easyEnergy (#133053)

homeassistant/components/easyenergy/__init__.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
@@ -10,10 +9,10 @@ from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.typing import ConfigType
 
 from .const import DOMAIN
-from .coordinator import EasyEnergyDataUpdateCoordinator
+from .coordinator import EasyEnergyConfigEntry, EasyEnergyDataUpdateCoordinator
 from .services import async_setup_services
 
-PLATFORMS = [Platform.SENSOR]
+PLATFORMS: list[Platform] = [Platform.SENSOR]
 CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
 
 
@@ -25,25 +24,22 @@ 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: EasyEnergyConfigEntry) -> bool:
     """Set up easyEnergy from a config entry."""
 
-    coordinator = EasyEnergyDataUpdateCoordinator(hass)
+    coordinator = EasyEnergyDataUpdateCoordinator(hass, entry)
     try:
         await coordinator.async_config_entry_first_refresh()
     except ConfigEntryNotReady:
         await coordinator.easyenergy.close()
         raise
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
 
     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: EasyEnergyConfigEntry) -> bool:
     """Unload easyEnergy 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/easyenergy/coordinator.py

@@ -21,6 +21,8 @@ from homeassistant.util import dt as dt_util
 
 from .const import DOMAIN, LOGGER, SCAN_INTERVAL, THRESHOLD_HOUR
 
+type EasyEnergyConfigEntry = ConfigEntry[EasyEnergyDataUpdateCoordinator]
+
 
 class EasyEnergyData(NamedTuple):
     """Class for defining data in dict."""
@@ -33,15 +35,16 @@ class EasyEnergyData(NamedTuple):
 class EasyEnergyDataUpdateCoordinator(DataUpdateCoordinator[EasyEnergyData]):
     """Class to manage fetching easyEnergy data from single endpoint."""
 
-    config_entry: ConfigEntry
+    config_entry: EasyEnergyConfigEntry
 
-    def __init__(self, hass: HomeAssistant) -> None:
+    def __init__(self, hass: HomeAssistant, entry: EasyEnergyConfigEntry) -> None:
         """Initialize global easyEnergy data updater."""
         super().__init__(
             hass,
             LOGGER,
             name=DOMAIN,
             update_interval=SCAN_INTERVAL,
+            config_entry=entry,
         )
 
         self.easyenergy = EasyEnergy(session=async_get_clientsession(hass))

homeassistant/components/easyenergy/diagnostics.py

@@ -5,12 +5,9 @@ from __future__ import annotations
 from datetime import timedelta
 from typing import Any
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from . import EasyEnergyDataUpdateCoordinator
-from .const import DOMAIN
-from .coordinator import EasyEnergyData
+from .coordinator import EasyEnergyConfigEntry, EasyEnergyData
 
 
 def get_gas_price(data: EasyEnergyData, hours: int) -> float | None:
@@ -32,41 +29,42 @@ def get_gas_price(data: EasyEnergyData, hours: int) -> float | None:
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: EasyEnergyConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: EasyEnergyDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator_data = entry.runtime_data.data
+    energy_today = coordinator_data.energy_today
 
     return {
         "entry": {
             "title": entry.title,
         },
         "energy_usage": {
-            "current_hour_price": coordinator.data.energy_today.current_usage_price,
-            "next_hour_price": coordinator.data.energy_today.price_at_time(
-                coordinator.data.energy_today.utcnow() + timedelta(hours=1)
+            "current_hour_price": energy_today.current_usage_price,
+            "next_hour_price": energy_today.price_at_time(
+                energy_today.utcnow() + timedelta(hours=1)
             ),
-            "average_price": coordinator.data.energy_today.average_usage_price,
-            "max_price": coordinator.data.energy_today.extreme_usage_prices[1],
-            "min_price": coordinator.data.energy_today.extreme_usage_prices[0],
-            "highest_price_time": coordinator.data.energy_today.highest_usage_price_time,
-            "lowest_price_time": coordinator.data.energy_today.lowest_usage_price_time,
-            "percentage_of_max": coordinator.data.energy_today.pct_of_max_usage,
+            "average_price": energy_today.average_usage_price,
+            "max_price": energy_today.extreme_usage_prices[1],
+            "min_price": energy_today.extreme_usage_prices[0],
+            "highest_price_time": energy_today.highest_usage_price_time,
+            "lowest_price_time": energy_today.lowest_usage_price_time,
+            "percentage_of_max": energy_today.pct_of_max_usage,
         },
         "energy_return": {
-            "current_hour_price": coordinator.data.energy_today.current_return_price,
-            "next_hour_price": coordinator.data.energy_today.price_at_time(
-                coordinator.data.energy_today.utcnow() + timedelta(hours=1), "return"
+            "current_hour_price": energy_today.current_return_price,
+            "next_hour_price": energy_today.price_at_time(
+                energy_today.utcnow() + timedelta(hours=1), "return"
             ),
-            "average_price": coordinator.data.energy_today.average_return_price,
-            "max_price": coordinator.data.energy_today.extreme_return_prices[1],
-            "min_price": coordinator.data.energy_today.extreme_return_prices[0],
-            "highest_price_time": coordinator.data.energy_today.highest_return_price_time,
-            "lowest_price_time": coordinator.data.energy_today.lowest_return_price_time,
-            "percentage_of_max": coordinator.data.energy_today.pct_of_max_return,
+            "average_price": energy_today.average_return_price,
+            "max_price": energy_today.extreme_return_prices[1],
+            "min_price": energy_today.extreme_return_prices[0],
+            "highest_price_time": energy_today.highest_return_price_time,
+            "lowest_price_time": energy_today.lowest_return_price_time,
+            "percentage_of_max": energy_today.pct_of_max_return,
         },
         "gas": {
-            "current_hour_price": get_gas_price(coordinator.data, 0),
-            "next_hour_price": get_gas_price(coordinator.data, 1),
+            "current_hour_price": get_gas_price(coordinator_data, 0),
+            "next_hour_price": get_gas_price(coordinator_data, 1),
         },
     }

homeassistant/components/easyenergy/sensor.py

@@ -13,7 +13,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CURRENCY_EURO,
     PERCENTAGE,
@@ -27,7 +26,11 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DOMAIN, SERVICE_TYPE_DEVICE_NAMES
-from .coordinator import EasyEnergyData, EasyEnergyDataUpdateCoordinator
+from .coordinator import (
+    EasyEnergyConfigEntry,
+    EasyEnergyData,
+    EasyEnergyDataUpdateCoordinator,
+)
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -208,10 +211,12 @@ def get_gas_price(data: EasyEnergyData, hours: int) -> float | None:
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: EasyEnergyConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up easyEnergy sensors based on a config entry."""
-    coordinator: EasyEnergyDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
     async_add_entities(
         EasyEnergySensorEntity(coordinator=coordinator, description=description)
         for description in SENSORS

homeassistant/components/easyenergy/services.py

@@ -10,7 +10,7 @@ from typing import Final
 from easyenergy import Electricity, Gas, VatOption
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigEntry, ConfigEntryState
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import (
     HomeAssistant,
     ServiceCall,
@@ -23,7 +23,7 @@ from homeassistant.helpers import selector
 from homeassistant.util import dt as dt_util
 
 from .const import DOMAIN
-from .coordinator import EasyEnergyDataUpdateCoordinator
+from .coordinator import EasyEnergyConfigEntry, EasyEnergyDataUpdateCoordinator
 
 ATTR_CONFIG_ENTRY: Final = "config_entry"
 ATTR_START: Final = "start"
@@ -91,7 +91,7 @@ def __get_coordinator(
 ) -> EasyEnergyDataUpdateCoordinator:
     """Get the coordinator from the entry."""
     entry_id: str = call.data[ATTR_CONFIG_ENTRY]
-    entry: ConfigEntry | None = hass.config_entries.async_get_entry(entry_id)
+    entry: EasyEnergyConfigEntry | None = hass.config_entries.async_get_entry(entry_id)
 
     if not entry:
         raise ServiceValidationError(
@@ -110,8 +110,7 @@ def __get_coordinator(
             },
         )
 
-    coordinator: EasyEnergyDataUpdateCoordinator = hass.data[DOMAIN][entry_id]
-    return coordinator
+    return entry.runtime_data
 
 
 async def __get_prices(

tests/components/easyenergy/test_init.py

@@ -4,7 +4,6 @@ from unittest.mock import MagicMock, patch
 
 from easyenergy import EasyEnergyConnectionError
 
-from homeassistant.components.easyenergy.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 
@@ -24,7 +23,6 @@ async def test_load_unload_config_entry(
     await hass.config_entries.async_unload(mock_config_entry.entry_id)
     await hass.async_block_till_done()
 
-    assert not hass.data.get(DOMAIN)
     assert mock_config_entry.state is ConfigEntryState.NOT_LOADED