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

pull/107510/head
Michael 2024-05-04 22:52:13 +02:00 committed by GitHub
parent 55ffc82be1
commit 4a25e67234
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 57 additions and 67 deletions

View File

@ -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:

View File

@ -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,

View File

@ -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"

View File

@ -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),

View File

@ -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,

View File

@ -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,
)

View File

@ -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,

View File

@ -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: