diff --git a/homeassistant/components/nextdns/__init__.py b/homeassistant/components/nextdns/__init__.py index c7e4a0842fb..f76e8755734 100644 --- a/homeassistant/components/nextdns/__init__.py +++ b/homeassistant/components/nextdns/__init__.py @@ -3,10 +3,21 @@ from __future__ import annotations import asyncio +from dataclasses import dataclass from datetime import timedelta from aiohttp.client_exceptions import ClientConnectorError -from nextdns import ApiError, NextDns +from nextdns import ( + AnalyticsDnssec, + AnalyticsEncryption, + AnalyticsIpVersions, + AnalyticsProtocols, + AnalyticsStatus, + ApiError, + ConnectionStatus, + NextDns, + Settings, +) from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_API_KEY, Platform @@ -23,7 +34,6 @@ from .const import ( ATTR_SETTINGS, ATTR_STATUS, CONF_PROFILE_ID, - DOMAIN, UPDATE_INTERVAL_ANALYTICS, UPDATE_INTERVAL_CONNECTION, UPDATE_INTERVAL_SETTINGS, @@ -39,6 +49,22 @@ from .coordinator import ( NextDnsUpdateCoordinator, ) +NextDnsConfigEntry = ConfigEntry["NextDnsData"] + + +@dataclass +class NextDnsData: + """Data for the NextDNS integration.""" + + connection: NextDnsUpdateCoordinator[ConnectionStatus] + dnssec: NextDnsUpdateCoordinator[AnalyticsDnssec] + encryption: NextDnsUpdateCoordinator[AnalyticsEncryption] + ip_versions: NextDnsUpdateCoordinator[AnalyticsIpVersions] + protocols: NextDnsUpdateCoordinator[AnalyticsProtocols] + settings: NextDnsUpdateCoordinator[Settings] + status: NextDnsUpdateCoordinator[AnalyticsStatus] + + PLATFORMS = [Platform.BINARY_SENSOR, Platform.BUTTON, Platform.SENSOR, Platform.SWITCH] COORDINATORS: list[tuple[str, type[NextDnsUpdateCoordinator], timedelta]] = [ (ATTR_CONNECTION, NextDnsConnectionUpdateCoordinator, UPDATE_INTERVAL_CONNECTION), @@ -51,7 +77,7 @@ COORDINATORS: list[tuple[str, type[NextDnsUpdateCoordinator], timedelta]] = [ ] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: NextDnsConfigEntry) -> bool: """Set up NextDNS as config entry.""" api_key = entry.data[CONF_API_KEY] profile_id = entry.data[CONF_PROFILE_ID] @@ -75,18 +101,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: await asyncio.gather(*tasks) - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinators + entry.runtime_data = NextDnsData(**coordinators) 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: NextDnsConfigEntry) -> bool: """Unload a config entry.""" - unload_ok: bool = 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) diff --git a/homeassistant/components/nextdns/binary_sensor.py b/homeassistant/components/nextdns/binary_sensor.py index c4ab58537cd..eaa5565bdc5 100644 --- a/homeassistant/components/nextdns/binary_sensor.py +++ b/homeassistant/components/nextdns/binary_sensor.py @@ -13,14 +13,13 @@ from homeassistant.components.binary_sensor import ( BinarySensorEntity, BinarySensorEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity -from .const import ATTR_CONNECTION, DOMAIN -from .coordinator import CoordinatorDataT, NextDnsConnectionUpdateCoordinator +from . import NextDnsConfigEntry +from .coordinator import CoordinatorDataT, NextDnsUpdateCoordinator PARALLEL_UPDATES = 1 @@ -54,13 +53,11 @@ SENSORS = ( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: NextDnsConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Add NextDNS entities from a config_entry.""" - coordinator: NextDnsConnectionUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][ - ATTR_CONNECTION - ] + coordinator = entry.runtime_data.connection async_add_entities( NextDnsBinarySensor(coordinator, description) for description in SENSORS @@ -68,7 +65,7 @@ async def async_setup_entry( class NextDnsBinarySensor( - CoordinatorEntity[NextDnsConnectionUpdateCoordinator], BinarySensorEntity + CoordinatorEntity[NextDnsUpdateCoordinator[ConnectionStatus]], BinarySensorEntity ): """Define an NextDNS binary sensor.""" @@ -77,7 +74,7 @@ class NextDnsBinarySensor( def __init__( self, - coordinator: NextDnsConnectionUpdateCoordinator, + coordinator: NextDnsUpdateCoordinator[ConnectionStatus], description: NextDnsBinarySensorEntityDescription, ) -> None: """Initialize.""" diff --git a/homeassistant/components/nextdns/button.py b/homeassistant/components/nextdns/button.py index d61c953f260..164d725b393 100644 --- a/homeassistant/components/nextdns/button.py +++ b/homeassistant/components/nextdns/button.py @@ -2,15 +2,16 @@ from __future__ import annotations +from nextdns import AnalyticsStatus + from homeassistant.components.button import ButtonEntity, ButtonEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity -from .const import ATTR_STATUS, DOMAIN -from .coordinator import NextDnsStatusUpdateCoordinator +from . import NextDnsConfigEntry +from .coordinator import NextDnsUpdateCoordinator PARALLEL_UPDATES = 1 @@ -22,27 +23,26 @@ CLEAR_LOGS_BUTTON = ButtonEntityDescription( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: NextDnsConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Add aNextDNS entities from a config_entry.""" - coordinator: NextDnsStatusUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][ - ATTR_STATUS - ] + coordinator = entry.runtime_data.status - buttons: list[NextDnsButton] = [] - buttons.append(NextDnsButton(coordinator, CLEAR_LOGS_BUTTON)) - - async_add_entities(buttons) + async_add_entities([NextDnsButton(coordinator, CLEAR_LOGS_BUTTON)]) -class NextDnsButton(CoordinatorEntity[NextDnsStatusUpdateCoordinator], ButtonEntity): +class NextDnsButton( + CoordinatorEntity[NextDnsUpdateCoordinator[AnalyticsStatus]], ButtonEntity +): """Define an NextDNS button.""" _attr_has_entity_name = True def __init__( self, - coordinator: NextDnsStatusUpdateCoordinator, + coordinator: NextDnsUpdateCoordinator[AnalyticsStatus], description: ButtonEntityDescription, ) -> None: """Initialize.""" diff --git a/homeassistant/components/nextdns/diagnostics.py b/homeassistant/components/nextdns/diagnostics.py index cade6476d82..31c0b7f0ca8 100644 --- a/homeassistant/components/nextdns/diagnostics.py +++ b/homeassistant/components/nextdns/diagnostics.py @@ -6,36 +6,25 @@ 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_API_KEY, CONF_UNIQUE_ID from homeassistant.core import HomeAssistant -from .const import ( - ATTR_DNSSEC, - ATTR_ENCRYPTION, - ATTR_IP_VERSIONS, - ATTR_PROTOCOLS, - ATTR_SETTINGS, - ATTR_STATUS, - CONF_PROFILE_ID, - DOMAIN, -) +from . import NextDnsConfigEntry +from .const import CONF_PROFILE_ID TO_REDACT = {CONF_API_KEY, CONF_PROFILE_ID, CONF_UNIQUE_ID} async def async_get_config_entry_diagnostics( - hass: HomeAssistant, config_entry: ConfigEntry + hass: HomeAssistant, config_entry: NextDnsConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - coordinators = hass.data[DOMAIN][config_entry.entry_id] - - dnssec_coordinator = coordinators[ATTR_DNSSEC] - encryption_coordinator = coordinators[ATTR_ENCRYPTION] - ip_versions_coordinator = coordinators[ATTR_IP_VERSIONS] - protocols_coordinator = coordinators[ATTR_PROTOCOLS] - settings_coordinator = coordinators[ATTR_SETTINGS] - status_coordinator = coordinators[ATTR_STATUS] + dnssec_coordinator = config_entry.runtime_data.dnssec + encryption_coordinator = config_entry.runtime_data.encryption + ip_versions_coordinator = config_entry.runtime_data.ip_versions + protocols_coordinator = config_entry.runtime_data.protocols + settings_coordinator = config_entry.runtime_data.settings + status_coordinator = config_entry.runtime_data.status return { "config_entry": async_redact_data(config_entry.as_dict(), TO_REDACT), diff --git a/homeassistant/components/nextdns/sensor.py b/homeassistant/components/nextdns/sensor.py index a034901aa41..b390ac93e06 100644 --- a/homeassistant/components/nextdns/sensor.py +++ b/homeassistant/components/nextdns/sensor.py @@ -19,20 +19,19 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import PERCENTAGE, EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity +from . import NextDnsConfigEntry from .const import ( ATTR_DNSSEC, ATTR_ENCRYPTION, ATTR_IP_VERSIONS, ATTR_PROTOCOLS, ATTR_STATUS, - DOMAIN, ) from .coordinator import CoordinatorDataT, NextDnsUpdateCoordinator @@ -301,14 +300,14 @@ SENSORS: tuple[NextDnsSensorEntityDescription, ...] = ( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: NextDnsConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Add a NextDNS entities from a config_entry.""" - coordinators = hass.data[DOMAIN][entry.entry_id] - async_add_entities( - NextDnsSensor(coordinators[description.coordinator_type], description) + NextDnsSensor( + getattr(entry.runtime_data, description.coordinator_type), description + ) for description in SENSORS ) diff --git a/homeassistant/components/nextdns/switch.py b/homeassistant/components/nextdns/switch.py index a6bbead131e..5e599d281d8 100644 --- a/homeassistant/components/nextdns/switch.py +++ b/homeassistant/components/nextdns/switch.py @@ -11,15 +11,14 @@ from aiohttp.client_exceptions import ClientConnectorError from nextdns import ApiError, Settings from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity -from .const import ATTR_SETTINGS, DOMAIN -from .coordinator import CoordinatorDataT, NextDnsSettingsUpdateCoordinator +from . import NextDnsConfigEntry +from .coordinator import CoordinatorDataT, NextDnsUpdateCoordinator PARALLEL_UPDATES = 1 @@ -526,19 +525,21 @@ SWITCHES = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: NextDnsConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Add NextDNS entities from a config_entry.""" - coordinator: NextDnsSettingsUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][ - ATTR_SETTINGS - ] + coordinator = entry.runtime_data.settings async_add_entities( NextDnsSwitch(coordinator, description) for description in SWITCHES ) -class NextDnsSwitch(CoordinatorEntity[NextDnsSettingsUpdateCoordinator], SwitchEntity): +class NextDnsSwitch( + CoordinatorEntity[NextDnsUpdateCoordinator[Settings]], SwitchEntity +): """Define an NextDNS switch.""" _attr_has_entity_name = True @@ -546,7 +547,7 @@ class NextDnsSwitch(CoordinatorEntity[NextDnsSettingsUpdateCoordinator], SwitchE def __init__( self, - coordinator: NextDnsSettingsUpdateCoordinator, + coordinator: NextDnsUpdateCoordinator[Settings], description: NextDnsSwitchEntityDescription, ) -> None: """Initialize."""