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 dataclasses import dataclass
from adguardhome import AdGuardHome, AdGuardHomeConnectionError
import voluptuous as vol
@ -24,7 +26,6 @@ from homeassistant.helpers.aiohttp_client import async_get_clientsession
from .const import (
CONF_FORCE,
DATA_ADGUARD_CLIENT,
DOMAIN,
SERVICE_ADD_URL,
SERVICE_DISABLE_URL,
@ -44,6 +45,14 @@ SERVICE_REFRESH_SCHEMA = vol.Schema(
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:
"""Set up AdGuard Home from a config entry."""
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,
)
hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {DATA_ADGUARD_CLIENT: adguard}
try:
await adguard.version()
version = await adguard.version()
except AdGuardHomeConnectionError as 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)
async def add_url(call: ServiceCall) -> None:

View File

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

View File

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

View File

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

View File

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