Store runtime data inside the config entry in Pi-Hole (#116806)

homeassistant/components/pi_hole/__init__.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 import logging
 
 from hole import Hole
@@ -28,13 +29,7 @@ from homeassistant.helpers.update_coordinator import (
     UpdateFailed,
 )
 
-from .const import (
+from .const import CONF_STATISTICS_ONLY, DOMAIN, MIN_TIME_BETWEEN_UPDATES
-    CONF_STATISTICS_ONLY,
-    DATA_KEY_API,
-    DATA_KEY_COORDINATOR,
-    DOMAIN,
-    MIN_TIME_BETWEEN_UPDATES,
-)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -47,8 +42,18 @@ PLATFORMS = [
     Platform.UPDATE,
 ]
 
+PiHoleConfigEntry = ConfigEntry["PiHoleData"]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+@dataclass
+class PiHoleData:
+    """Runtime data definition."""
+
+    api: Hole
+    coordinator: DataUpdateCoordinator[None]
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: PiHoleConfigEntry) -> bool:
     """Set up Pi-hole entry."""
     name = entry.data[CONF_NAME]
     host = entry.data[CONF_HOST]
@@ -126,11 +131,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})
+    entry.runtime_data = PiHoleData(api, coordinator)
-    hass.data[DOMAIN][entry.entry_id] = {
-        DATA_KEY_API: api,
-        DATA_KEY_COORDINATOR: coordinator,
-    }
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -139,19 +140,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload Pi-hole 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
 
 
-class PiHoleEntity(CoordinatorEntity):
+class PiHoleEntity(CoordinatorEntity[DataUpdateCoordinator[None]]):
     """Representation of a Pi-hole entity."""
 
     def __init__(
         self,
         api: Hole,
-        coordinator: DataUpdateCoordinator,
+        coordinator: DataUpdateCoordinator[None],
         name: str,
         server_unique_id: str,
     ) -> None:

homeassistant/components/pi_hole/binary_sensor.py

@@ -12,14 +12,12 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from . import PiHoleEntity
+from . import PiHoleConfigEntry, PiHoleEntity
-from .const import DATA_KEY_API, DATA_KEY_COORDINATOR, DOMAIN as PIHOLE_DOMAIN
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -40,16 +38,18 @@ BINARY_SENSOR_TYPES: tuple[PiHoleBinarySensorEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: PiHoleConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Pi-hole binary sensor."""
     name = entry.data[CONF_NAME]
-    hole_data = hass.data[PIHOLE_DOMAIN][entry.entry_id]
+    hole_data = entry.runtime_data
 
     binary_sensors = [
         PiHoleBinarySensor(
-            hole_data[DATA_KEY_API],
+            hole_data.api,
-            hole_data[DATA_KEY_COORDINATOR],
+            hole_data.coordinator,
             name,
             entry.entry_id,
             description,
@@ -69,7 +69,7 @@ class PiHoleBinarySensor(PiHoleEntity, BinarySensorEntity):
     def __init__(
         self,
         api: Hole,
-        coordinator: DataUpdateCoordinator,
+        coordinator: DataUpdateCoordinator[None],
         name: str,
         server_unique_id: str,
         description: PiHoleBinarySensorEntityDescription,

homeassistant/components/pi_hole/const.py

@@ -17,6 +17,3 @@ SERVICE_DISABLE = "disable"
 SERVICE_DISABLE_ATTR_DURATION = "duration"
 
 MIN_TIME_BETWEEN_UPDATES = timedelta(minutes=5)
-
-DATA_KEY_API = "api"
-DATA_KEY_COORDINATOR = "coordinator"

homeassistant/components/pi_hole/diagnostics.py

@@ -4,23 +4,20 @@ from __future__ import annotations
 
 from typing import Any
 
-from hole import Hole
-
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY
 from homeassistant.core import HomeAssistant
 
-from .const import DATA_KEY_API, DOMAIN
+from . import PiHoleConfigEntry
 
 TO_REDACT = {CONF_API_KEY}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: PiHoleConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    api: Hole = hass.data[DOMAIN][entry.entry_id][DATA_KEY_API]
+    api = entry.runtime_data.api
 
     return {
         "entry": async_redact_data(entry.as_dict(), TO_REDACT),

homeassistant/components/pi_hole/sensor.py

@@ -5,15 +5,13 @@ from __future__ import annotations
 from hole import Hole
 
 from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME, PERCENTAGE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from . import PiHoleEntity
+from . import PiHoleConfigEntry, PiHoleEntity
-from .const import DATA_KEY_API, DATA_KEY_COORDINATOR, DOMAIN as PIHOLE_DOMAIN
 
 SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
     SensorEntityDescription(
@@ -65,15 +63,17 @@ SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: PiHoleConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Pi-hole sensor."""
     name = entry.data[CONF_NAME]
-    hole_data = hass.data[PIHOLE_DOMAIN][entry.entry_id]
+    hole_data = entry.runtime_data
     sensors = [
         PiHoleSensor(
-            hole_data[DATA_KEY_API],
+            hole_data.api,
-            hole_data[DATA_KEY_COORDINATOR],
+            hole_data.coordinator,
             name,
             entry.entry_id,
             description,
@@ -92,7 +92,7 @@ class PiHoleSensor(PiHoleEntity, SensorEntity):
     def __init__(
         self,
         api: Hole,
-        coordinator: DataUpdateCoordinator,
+        coordinator: DataUpdateCoordinator[None],
         name: str,
         server_unique_id: str,
         description: SensorEntityDescription,

homeassistant/components/pi_hole/switch.py

@@ -9,34 +9,29 @@ from hole.exceptions import HoleError
 import voluptuous as vol
 
 from homeassistant.components.switch import SwitchEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation as cv, entity_platform
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import PiHoleEntity
+from . import PiHoleConfigEntry, PiHoleEntity
-from .const import (
+from .const import SERVICE_DISABLE, SERVICE_DISABLE_ATTR_DURATION
-    DATA_KEY_API,
-    DATA_KEY_COORDINATOR,
-    DOMAIN as PIHOLE_DOMAIN,
-    SERVICE_DISABLE,
-    SERVICE_DISABLE_ATTR_DURATION,
-)
 
 _LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: PiHoleConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Pi-hole switch."""
     name = entry.data[CONF_NAME]
-    hole_data = hass.data[PIHOLE_DOMAIN][entry.entry_id]
+    hole_data = entry.runtime_data
     switches = [
         PiHoleSwitch(
-            hole_data[DATA_KEY_API],
+            hole_data.api,
-            hole_data[DATA_KEY_COORDINATOR],
+            hole_data.coordinator,
             name,
             entry.entry_id,
         )

homeassistant/components/pi_hole/update.py

@@ -8,14 +8,12 @@ from dataclasses import dataclass
 from hole import Hole
 
 from homeassistant.components.update import UpdateEntity, UpdateEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME, EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from . import PiHoleEntity
+from . import PiHoleConfigEntry, PiHoleEntity
-from .const import DATA_KEY_API, DATA_KEY_COORDINATOR, DOMAIN
 
 
 @dataclass(frozen=True)
@@ -60,16 +58,18 @@ UPDATE_ENTITY_TYPES: tuple[PiHoleUpdateEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: PiHoleConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Pi-hole update entities."""
     name = entry.data[CONF_NAME]
-    hole_data = hass.data[DOMAIN][entry.entry_id]
+    hole_data = entry.runtime_data
 
     async_add_entities(
         PiHoleUpdateEntity(
-            hole_data[DATA_KEY_API],
+            hole_data.api,
-            hole_data[DATA_KEY_COORDINATOR],
+            hole_data.coordinator,
             name,
             entry.entry_id,
             description,
@@ -87,7 +87,7 @@ class PiHoleUpdateEntity(PiHoleEntity, UpdateEntity):
     def __init__(
         self,
         api: Hole,
-        coordinator: DataUpdateCoordinator,
+        coordinator: DataUpdateCoordinator[None],
         name: str,
         server_unique_id: str,
         description: PiHoleUpdateEntityDescription,

tests/components/pi_hole/test_init.py

@@ -7,11 +7,13 @@ from hole.exceptions import HoleError
 import pytest
 
 from homeassistant.components import pi_hole, switch
+from homeassistant.components.pi_hole import PiHoleData
 from homeassistant.components.pi_hole.const import (
     CONF_STATISTICS_ONLY,
     SERVICE_DISABLE,
     SERVICE_DISABLE_ATTR_DURATION,
 )
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import ATTR_ENTITY_ID, CONF_HOST, CONF_NAME
 from homeassistant.core import HomeAssistant
 
@@ -182,12 +184,13 @@ async def test_unload(hass: HomeAssistant) -> None:
     with _patch_init_hole(mocked_hole):
         await hass.config_entries.async_setup(entry.entry_id)
         await hass.async_block_till_done()
-    assert entry.entry_id in hass.data[pi_hole.DOMAIN]
+    assert entry.state is ConfigEntryState.LOADED
+    assert isinstance(entry.runtime_data, PiHoleData)
     assert await hass.config_entries.async_unload(entry.entry_id)
 
     await hass.async_block_till_done()
 
-    assert entry.entry_id not in hass.data[pi_hole.DOMAIN]
+    assert entry.state is ConfigEntryState.NOT_LOADED
 
 
 async def test_remove_obsolete(hass: HomeAssistant) -> None: