Use ConfigEntry runtime_data in Discovergy (#116671)

homeassistant/components/discovergy/__init__.py

@@ -12,16 +12,15 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers.httpx_client import get_async_client
 
-from .const import DOMAIN
 from .coordinator import DiscovergyUpdateCoordinator
 
 PLATFORMS = [Platform.SENSOR]
 
+DiscovergyConfigEntry = ConfigEntry[list[DiscovergyUpdateCoordinator]]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: DiscovergyConfigEntry) -> bool:
     """Set up Discovergy from a config entry."""
-    hass.data.setdefault(DOMAIN, {})
-
     client = Discovergy(
         email=entry.data[CONF_EMAIL],
         password=entry.data[CONF_PASSWORD],
@@ -53,7 +52,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         await coordinator.async_config_entry_first_refresh()
         coordinators.append(coordinator)
 
-    hass.data[DOMAIN][entry.entry_id] = coordinators
+    entry.runtime_data = coordinators
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     entry.async_on_unload(entry.add_update_listener(async_reload_entry))
@@ -63,11 +62,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
 
 
 async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:

homeassistant/components/discovergy/diagnostics.py

@@ -6,11 +6,9 @@ from dataclasses import asdict
 from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
+from . import DiscovergyConfigEntry
-from .coordinator import DiscovergyUpdateCoordinator
 
 TO_REDACT_METER = {
     "serial_number",
@@ -22,14 +20,13 @@ TO_REDACT_METER = {
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: DiscovergyConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
     flattened_meter: list[dict] = []
     last_readings: dict[str, dict] = {}
-    coordinators: list[DiscovergyUpdateCoordinator] = hass.data[DOMAIN][entry.entry_id]
 
-    for coordinator in coordinators:
+    for coordinator in entry.runtime_data:
         # make a dict of meter data and redact some data
         flattened_meter.append(
             async_redact_data(asdict(coordinator.meter), TO_REDACT_METER)

homeassistant/components/discovergy/sensor.py

@@ -12,7 +12,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     EntityCategory,
     UnitOfElectricPotential,
@@ -25,6 +24,7 @@ from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from . import DiscovergyConfigEntry
 from .const import DOMAIN, MANUFACTURER
 from .coordinator import DiscovergyUpdateCoordinator
 
@@ -163,13 +163,13 @@ ADDITIONAL_SENSORS: tuple[DiscovergySensorEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: DiscovergyConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Discovergy sensors."""
-    coordinators: list[DiscovergyUpdateCoordinator] = hass.data[DOMAIN][entry.entry_id]
-
     entities: list[DiscovergySensor] = []
-    for coordinator in coordinators:
+    for coordinator in entry.runtime_data:
         sensors: tuple[DiscovergySensorEntityDescription, ...] = ()
 
         # select sensor descriptions based on meter type and combine with additional sensors

tests/components/discovergy/conftest.py

@@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, patch
 from pydiscovergy.models import Reading
 import pytest
 
-from homeassistant.components.discovergy import DOMAIN
+from homeassistant.components.discovergy.const import DOMAIN
 from homeassistant.const import CONF_EMAIL, CONF_PASSWORD
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component