diff --git a/homeassistant/components/pi_hole/__init__.py b/homeassistant/components/pi_hole/__init__.py index f892114b26c..05d301b5250 100644 --- a/homeassistant/components/pi_hole/__init__.py +++ b/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 ( - CONF_STATISTICS_ONLY, - DATA_KEY_API, - DATA_KEY_COORDINATOR, - DOMAIN, - MIN_TIME_BETWEEN_UPDATES, -) +from .const import CONF_STATISTICS_ONLY, 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, {}) - hass.data[DOMAIN][entry.entry_id] = { - DATA_KEY_API: api, - DATA_KEY_COORDINATOR: coordinator, - } + entry.runtime_data = PiHoleData(api, 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) - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) -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: diff --git a/homeassistant/components/pi_hole/binary_sensor.py b/homeassistant/components/pi_hole/binary_sensor.py index 0593d12faa7..001a2ebcee8 100644 --- a/homeassistant/components/pi_hole/binary_sensor.py +++ b/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 .const import DATA_KEY_API, DATA_KEY_COORDINATOR, DOMAIN as PIHOLE_DOMAIN +from . import PiHoleConfigEntry, PiHoleEntity @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[DATA_KEY_COORDINATOR], + hole_data.api, + 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, diff --git a/homeassistant/components/pi_hole/const.py b/homeassistant/components/pi_hole/const.py index b6c97bc6118..c81e6504dff 100644 --- a/homeassistant/components/pi_hole/const.py +++ b/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" diff --git a/homeassistant/components/pi_hole/diagnostics.py b/homeassistant/components/pi_hole/diagnostics.py index 46efebaf475..115c04c8234 100644 --- a/homeassistant/components/pi_hole/diagnostics.py +++ b/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), diff --git a/homeassistant/components/pi_hole/sensor.py b/homeassistant/components/pi_hole/sensor.py index a62252d10c1..14ad3ac82dd 100644 --- a/homeassistant/components/pi_hole/sensor.py +++ b/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 .const import DATA_KEY_API, DATA_KEY_COORDINATOR, DOMAIN as PIHOLE_DOMAIN +from . import PiHoleConfigEntry, PiHoleEntity 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[DATA_KEY_COORDINATOR], + hole_data.api, + 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, diff --git a/homeassistant/components/pi_hole/switch.py b/homeassistant/components/pi_hole/switch.py index 963ee7c9738..83ed3e6d787 100644 --- a/homeassistant/components/pi_hole/switch.py +++ b/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 .const import ( - DATA_KEY_API, - DATA_KEY_COORDINATOR, - DOMAIN as PIHOLE_DOMAIN, - SERVICE_DISABLE, - SERVICE_DISABLE_ATTR_DURATION, -) +from . import PiHoleConfigEntry, PiHoleEntity +from .const import 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[DATA_KEY_COORDINATOR], + hole_data.api, + hole_data.coordinator, name, entry.entry_id, ) diff --git a/homeassistant/components/pi_hole/update.py b/homeassistant/components/pi_hole/update.py index 75d4f91f2be..db78d3ab0a5 100644 --- a/homeassistant/components/pi_hole/update.py +++ b/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 .const import DATA_KEY_API, DATA_KEY_COORDINATOR, DOMAIN +from . import PiHoleConfigEntry, PiHoleEntity @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[DATA_KEY_COORDINATOR], + hole_data.api, + 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, diff --git a/tests/components/pi_hole/test_init.py b/tests/components/pi_hole/test_init.py index a58a46680bb..3c8f66a82d0 100644 --- a/tests/components/pi_hole/test_init.py +++ b/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: