Add dataclass to store AdGuard data (#115668)

* Add dataclass to store AdGuard data

* Unify version call

homeassistant/components/adguard/__init__.py

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

homeassistant/components/adguard/const.py

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

homeassistant/components/adguard/entity.py

@@ -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(
+            sw_version=self.data.version,
-                DATA_ADGUARD_VERSION
-            ),
             configuration_url=config_url,
         )

homeassistant/components/adguard/sensor.py

@@ -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]
+    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(
-        [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,
+                self.adguard.host,
-                str(adguard.port),
+                str(self.adguard.port),
                 "sensor",
                 description.key,
             ]

homeassistant/components/adguard/switch.py

@@ -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]
+    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(
-        [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: