Add dataclass to store AdGuard data (#115668)

* Add dataclass to store AdGuard data

* Unify version call
This commit is contained in:
Marc Mueller 2024-04-16 08:02:27 +02:00 committed by GitHub
parent 11ff00f637
commit 6a7a44c998
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 43 additions and 45 deletions

View File

@ -2,6 +2,8 @@
from __future__ import annotations from __future__ import annotations
from dataclasses import dataclass
from adguardhome import AdGuardHome, AdGuardHomeConnectionError from adguardhome import AdGuardHome, AdGuardHomeConnectionError
import voluptuous as vol import voluptuous as vol
@ -24,7 +26,6 @@ from homeassistant.helpers.aiohttp_client import async_get_clientsession
from .const import ( from .const import (
CONF_FORCE, CONF_FORCE,
DATA_ADGUARD_CLIENT,
DOMAIN, DOMAIN,
SERVICE_ADD_URL, SERVICE_ADD_URL,
SERVICE_DISABLE_URL, SERVICE_DISABLE_URL,
@ -44,6 +45,14 @@ SERVICE_REFRESH_SCHEMA = vol.Schema(
PLATFORMS = [Platform.SENSOR, Platform.SWITCH] PLATFORMS = [Platform.SENSOR, Platform.SWITCH]
@dataclass
class AdGuardData:
"""Adguard data type."""
client: AdGuardHome
version: str
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up AdGuard Home from a config entry.""" """Set up AdGuard Home from a config entry."""
session = async_get_clientsession(hass, entry.data[CONF_VERIFY_SSL]) session = async_get_clientsession(hass, entry.data[CONF_VERIFY_SSL])
@ -57,13 +66,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
session=session, session=session,
) )
hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {DATA_ADGUARD_CLIENT: adguard}
try: try:
await adguard.version() version = await adguard.version()
except AdGuardHomeConnectionError as exception: except AdGuardHomeConnectionError as exception:
raise ConfigEntryNotReady from exception raise ConfigEntryNotReady from exception
hass.data.setdefault(DOMAIN, {})[entry.entry_id] = AdGuardData(adguard, version)
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
async def add_url(call: ServiceCall) -> None: async def add_url(call: ServiceCall) -> None:

View File

@ -6,9 +6,6 @@ DOMAIN = "adguard"
LOGGER = logging.getLogger(__package__) LOGGER = logging.getLogger(__package__)
DATA_ADGUARD_CLIENT = "adguard_client"
DATA_ADGUARD_VERSION = "adguard_version"
CONF_FORCE = "force" CONF_FORCE = "force"
SERVICE_ADD_URL = "add_url" SERVICE_ADD_URL = "add_url"

View File

@ -2,13 +2,14 @@
from __future__ import annotations from __future__ import annotations
from adguardhome import AdGuardHome, AdGuardHomeError from adguardhome import AdGuardHomeError
from homeassistant.config_entries import SOURCE_HASSIO, ConfigEntry from homeassistant.config_entries import SOURCE_HASSIO, ConfigEntry
from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
from homeassistant.helpers.entity import Entity from homeassistant.helpers.entity import Entity
from .const import DATA_ADGUARD_VERSION, DOMAIN, LOGGER from . import AdGuardData
from .const import DOMAIN, LOGGER
class AdGuardHomeEntity(Entity): class AdGuardHomeEntity(Entity):
@ -19,12 +20,13 @@ class AdGuardHomeEntity(Entity):
def __init__( def __init__(
self, self,
adguard: AdGuardHome, data: AdGuardData,
entry: ConfigEntry, entry: ConfigEntry,
) -> None: ) -> None:
"""Initialize the AdGuard Home entity.""" """Initialize the AdGuard Home entity."""
self._entry = entry self._entry = entry
self.adguard = adguard self.data = data
self.adguard = data.client
async def async_update(self) -> None: async def async_update(self) -> None:
"""Update AdGuard Home entity.""" """Update AdGuard Home entity."""
@ -68,8 +70,6 @@ class AdGuardHomeEntity(Entity):
}, },
manufacturer="AdGuard Team", manufacturer="AdGuard Team",
name="AdGuard Home", name="AdGuard Home",
sw_version=self.hass.data[DOMAIN][self._entry.entry_id].get( sw_version=self.data.version,
DATA_ADGUARD_VERSION
),
configuration_url=config_url, configuration_url=config_url,
) )

View File

@ -7,16 +7,16 @@ from dataclasses import dataclass
from datetime import timedelta from datetime import timedelta
from typing import Any from typing import Any
from adguardhome import AdGuardHome, AdGuardHomeConnectionError from adguardhome import AdGuardHome
from homeassistant.components.sensor import SensorEntity, SensorEntityDescription from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.const import PERCENTAGE, UnitOfTime from homeassistant.const import PERCENTAGE, UnitOfTime
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.exceptions import PlatformNotReady
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .const import DATA_ADGUARD_CLIENT, DATA_ADGUARD_VERSION, DOMAIN from . import AdGuardData
from .const import DOMAIN
from .entity import AdGuardHomeEntity from .entity import AdGuardHomeEntity
SCAN_INTERVAL = timedelta(seconds=300) SCAN_INTERVAL = timedelta(seconds=300)
@ -89,17 +89,10 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up AdGuard Home sensor based on a config entry.""" """Set up AdGuard Home sensor based on a config entry."""
adguard = hass.data[DOMAIN][entry.entry_id][DATA_ADGUARD_CLIENT] data: AdGuardData = hass.data[DOMAIN][entry.entry_id]
try:
version = await adguard.version()
except AdGuardHomeConnectionError as exception:
raise PlatformNotReady from exception
hass.data[DOMAIN][entry.entry_id][DATA_ADGUARD_VERSION] = version
async_add_entities( async_add_entities(
[AdGuardHomeSensor(adguard, entry, description) for description in SENSORS], [AdGuardHomeSensor(data, entry, description) for description in SENSORS],
True, True,
) )
@ -111,18 +104,18 @@ class AdGuardHomeSensor(AdGuardHomeEntity, SensorEntity):
def __init__( def __init__(
self, self,
adguard: AdGuardHome, data: AdGuardData,
entry: ConfigEntry, entry: ConfigEntry,
description: AdGuardHomeEntityDescription, description: AdGuardHomeEntityDescription,
) -> None: ) -> None:
"""Initialize AdGuard Home sensor.""" """Initialize AdGuard Home sensor."""
super().__init__(adguard, entry) super().__init__(data, entry)
self.entity_description = description self.entity_description = description
self._attr_unique_id = "_".join( self._attr_unique_id = "_".join(
[ [
DOMAIN, DOMAIN,
adguard.host, self.adguard.host,
str(adguard.port), str(self.adguard.port),
"sensor", "sensor",
description.key, description.key,
] ]

View File

@ -7,15 +7,15 @@ from dataclasses import dataclass
from datetime import timedelta from datetime import timedelta
from typing import Any from typing import Any
from adguardhome import AdGuardHome, AdGuardHomeConnectionError, AdGuardHomeError from adguardhome import AdGuardHome, AdGuardHomeError
from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.exceptions import PlatformNotReady
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .const import DATA_ADGUARD_CLIENT, DATA_ADGUARD_VERSION, DOMAIN, LOGGER from . import AdGuardData
from .const import DOMAIN, LOGGER
from .entity import AdGuardHomeEntity from .entity import AdGuardHomeEntity
SCAN_INTERVAL = timedelta(seconds=10) SCAN_INTERVAL = timedelta(seconds=10)
@ -83,17 +83,10 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up AdGuard Home switch based on a config entry.""" """Set up AdGuard Home switch based on a config entry."""
adguard = hass.data[DOMAIN][entry.entry_id][DATA_ADGUARD_CLIENT] data: AdGuardData = hass.data[DOMAIN][entry.entry_id]
try:
version = await adguard.version()
except AdGuardHomeConnectionError as exception:
raise PlatformNotReady from exception
hass.data[DOMAIN][entry.entry_id][DATA_ADGUARD_VERSION] = version
async_add_entities( async_add_entities(
[AdGuardHomeSwitch(adguard, entry, description) for description in SWITCHES], [AdGuardHomeSwitch(data, entry, description) for description in SWITCHES],
True, True,
) )
@ -105,15 +98,21 @@ class AdGuardHomeSwitch(AdGuardHomeEntity, SwitchEntity):
def __init__( def __init__(
self, self,
adguard: AdGuardHome, data: AdGuardData,
entry: ConfigEntry, entry: ConfigEntry,
description: AdGuardHomeSwitchEntityDescription, description: AdGuardHomeSwitchEntityDescription,
) -> None: ) -> None:
"""Initialize AdGuard Home switch.""" """Initialize AdGuard Home switch."""
super().__init__(adguard, entry) super().__init__(data, entry)
self.entity_description = description self.entity_description = description
self._attr_unique_id = "_".join( self._attr_unique_id = "_".join(
[DOMAIN, adguard.host, str(adguard.port), "switch", description.key] [
DOMAIN,
self.adguard.host,
str(self.adguard.port),
"switch",
description.key,
]
) )
async def async_turn_off(self, **kwargs: Any) -> None: async def async_turn_off(self, **kwargs: Any) -> None: