Use ConfigEntry runtime_data in Pure Energie (#131061)

homeassistant/components/pure_energie/__init__.py

@@ -7,13 +7,14 @@ from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 
-from .const import DOMAIN
 from .coordinator import PureEnergieDataUpdateCoordinator
 
-PLATFORMS = [Platform.SENSOR]
+PLATFORMS: list[Platform] = [Platform.SENSOR]
+
+type PureEnergieConfigEntry = ConfigEntry[PureEnergieDataUpdateCoordinator]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: PureEnergieConfigEntry) -> bool:
     """Set up Pure Energie from a config entry."""
 
     coordinator = PureEnergieDataUpdateCoordinator(hass)
@@ -23,14 +24,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         await coordinator.gridnet.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: PureEnergieConfigEntry
+) -> bool:
     """Unload Pure Energie config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        del hass.data[DOMAIN][entry.entry_id]
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/pure_energie/diagnostics.py

@@ -6,12 +6,10 @@ from dataclasses import asdict
 from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
-from .coordinator import PureEnergieDataUpdateCoordinator
+from . import PureEnergieConfigEntry
 
 TO_REDACT = {
     CONF_HOST,
@@ -20,18 +18,18 @@ TO_REDACT = {
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: PureEnergieConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: PureEnergieDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-
     return {
         "entry": {
             "title": entry.title,
             "data": async_redact_data(entry.data, TO_REDACT),
         },
         "data": {
-            "device": async_redact_data(asdict(coordinator.data.device), TO_REDACT),
-            "smartbridge": asdict(coordinator.data.smartbridge),
+            "device": async_redact_data(
+                asdict(entry.runtime_data.data.device), TO_REDACT
+            ),
+            "smartbridge": asdict(entry.runtime_data.data.smartbridge),
         },
     }

homeassistant/components/pure_energie/sensor.py

@@ -12,13 +12,13 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, UnitOfEnergy, UnitOfPower
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from . import PureEnergieConfigEntry
 from .const import DOMAIN
 from .coordinator import PureEnergieData, PureEnergieDataUpdateCoordinator
 
@@ -59,12 +59,13 @@ SENSORS: tuple[PureEnergieSensorEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: PureEnergieConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Pure Energie Sensors based on a config entry."""
     async_add_entities(
         PureEnergieSensorEntity(
-            coordinator=hass.data[DOMAIN][entry.entry_id],
             description=description,
             entry=entry,
         )
@@ -83,21 +84,22 @@ class PureEnergieSensorEntity(
     def __init__(
         self,
         *,
-        coordinator: PureEnergieDataUpdateCoordinator,
         description: PureEnergieSensorEntityDescription,
-        entry: ConfigEntry,
+        entry: PureEnergieConfigEntry,
     ) -> None:
         """Initialize Pure Energie sensor."""
-        super().__init__(coordinator=coordinator)
+        super().__init__(coordinator=entry.runtime_data)
         self.entity_id = f"{SENSOR_DOMAIN}.pem_{description.key}"
         self.entity_description = description
-        self._attr_unique_id = f"{coordinator.data.device.n2g_id}_{description.key}"
+        self._attr_unique_id = (
+            f"{entry.runtime_data.data.device.n2g_id}_{description.key}"
+        )
         self._attr_device_info = DeviceInfo(
-            identifiers={(DOMAIN, coordinator.data.device.n2g_id)},
-            configuration_url=f"http://{coordinator.config_entry.data[CONF_HOST]}",
-            sw_version=coordinator.data.device.firmware,
-            manufacturer=coordinator.data.device.manufacturer,
-            model=coordinator.data.device.model,
+            identifiers={(DOMAIN, entry.runtime_data.data.device.n2g_id)},
+            configuration_url=f"http://{entry.runtime_data.config_entry.data[CONF_HOST]}",
+            sw_version=entry.runtime_data.data.device.firmware,
+            manufacturer=entry.runtime_data.data.device.manufacturer,
+            model=entry.runtime_data.data.device.model,
             name=entry.title,
         )
 

tests/components/pure_energie/test_init.py

@@ -5,7 +5,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
 from gridnet import GridNetConnectionError
 import pytest
 
-from homeassistant.components.pure_energie.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 
@@ -32,7 +31,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