Add entity descriptions in AdGuard Home switches (#75229)

homeassistant/components/adguard/switch.py

@@ -1,13 +1,15 @@
 """Support for AdGuard Home switches."""
 from __future__ import annotations
 
+from collections.abc import Callable, Coroutine
+from dataclasses import dataclass
 from datetime import timedelta
 import logging
 from typing import Any
 
 from adguardhome import AdGuardHome, AdGuardHomeConnectionError, AdGuardHomeError
 
-from homeassistant.components.switch import SwitchEntity
+from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import PlatformNotReady
@@ -22,6 +24,74 @@ SCAN_INTERVAL = timedelta(seconds=10)
 PARALLEL_UPDATES = 1
 
 
+@dataclass
+class AdGuardHomeSwitchEntityDescriptionMixin:
+    """Mixin for required keys."""
+
+    is_on_fn: Callable[[AdGuardHome], Callable[[], Coroutine[Any, Any, bool]]]
+    turn_on_fn: Callable[[AdGuardHome], Callable[[], Coroutine[Any, Any, None]]]
+    turn_off_fn: Callable[[AdGuardHome], Callable[[], Coroutine[Any, Any, None]]]
+
+
+@dataclass
+class AdGuardHomeSwitchEntityDescription(
+    SwitchEntityDescription, AdGuardHomeSwitchEntityDescriptionMixin
+):
+    """Describes AdGuard Home switch entity."""
+
+
+SWITCHES: tuple[AdGuardHomeSwitchEntityDescription, ...] = (
+    AdGuardHomeSwitchEntityDescription(
+        key="protection",
+        name="Protection",
+        icon="mdi:shield-check",
+        is_on_fn=lambda adguard: adguard.protection_enabled,
+        turn_on_fn=lambda adguard: adguard.enable_protection,
+        turn_off_fn=lambda adguard: adguard.disable_protection,
+    ),
+    AdGuardHomeSwitchEntityDescription(
+        key="parental",
+        name="Parental control",
+        icon="mdi:shield-check",
+        is_on_fn=lambda adguard: adguard.parental.enabled,
+        turn_on_fn=lambda adguard: adguard.parental.enable,
+        turn_off_fn=lambda adguard: adguard.parental.disable,
+    ),
+    AdGuardHomeSwitchEntityDescription(
+        key="safesearch",
+        name="Safe search",
+        icon="mdi:shield-check",
+        is_on_fn=lambda adguard: adguard.safesearch.enabled,
+        turn_on_fn=lambda adguard: adguard.safesearch.enable,
+        turn_off_fn=lambda adguard: adguard.safesearch.disable,
+    ),
+    AdGuardHomeSwitchEntityDescription(
+        key="safebrowsing",
+        name="Safe browsing",
+        icon="mdi:shield-check",
+        is_on_fn=lambda adguard: adguard.safebrowsing.enabled,
+        turn_on_fn=lambda adguard: adguard.safebrowsing.enable,
+        turn_off_fn=lambda adguard: adguard.safebrowsing.disable,
+    ),
+    AdGuardHomeSwitchEntityDescription(
+        key="filtering",
+        name="Filtering",
+        icon="mdi:shield-check",
+        is_on_fn=lambda adguard: adguard.filtering.enabled,
+        turn_on_fn=lambda adguard: adguard.filtering.enable,
+        turn_off_fn=lambda adguard: adguard.filtering.disable,
+    ),
+    AdGuardHomeSwitchEntityDescription(
+        key="querylog",
+        name="Query log",
+        icon="mdi:shield-check",
+        is_on_fn=lambda adguard: adguard.querylog.enabled,
+        turn_on_fn=lambda adguard: adguard.querylog.enable,
+        turn_off_fn=lambda adguard: adguard.querylog.disable,
+    ),
+)
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     entry: ConfigEntry,
@@ -37,199 +107,54 @@ async def async_setup_entry(
 
     hass.data[DOMAIN][entry.entry_id][DATA_ADGUARD_VERSION] = version
 
-    switches = [
+    async_add_entities(
-        AdGuardHomeProtectionSwitch(adguard, entry),
+        [AdGuardHomeSwitch(adguard, entry, description) for description in SWITCHES],
-        AdGuardHomeFilteringSwitch(adguard, entry),
+        True,
-        AdGuardHomeParentalSwitch(adguard, entry),
+    )
-        AdGuardHomeSafeBrowsingSwitch(adguard, entry),
-        AdGuardHomeSafeSearchSwitch(adguard, entry),
-        AdGuardHomeQueryLogSwitch(adguard, entry),
-    ]
-    async_add_entities(switches, True)
 
 
 class AdGuardHomeSwitch(AdGuardHomeDeviceEntity, SwitchEntity):
     """Defines a AdGuard Home switch."""
 
+    entity_description: AdGuardHomeSwitchEntityDescription
+
     def __init__(
         self,
         adguard: AdGuardHome,
         entry: ConfigEntry,
-        name: str,
+        description: AdGuardHomeSwitchEntityDescription,
-        icon: str,
-        key: str,
-        enabled_default: bool = True,
     ) -> None:
         """Initialize AdGuard Home switch."""
-        self._state = False
+        self.entity_description = description
-        self._key = key
-        super().__init__(adguard, entry, name, icon, enabled_default)
 
-    @property
+        self._attr_unique_id = "_".join(
-    def unique_id(self) -> str:
+            [DOMAIN, adguard.host, str(adguard.port), "switch", description.key]
-        """Return the unique ID for this sensor."""
-        return "_".join(
-            [DOMAIN, self.adguard.host, str(self.adguard.port), "switch", self._key]
         )
 
-    @property
+        super().__init__(
-    def is_on(self) -> bool:
+            adguard,
-        """Return the state of the switch."""
+            entry,
-        return self._state
+            description.name,
+            description.icon,
+            description.entity_registry_enabled_default,
+        )
 
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off the switch."""
         try:
-            await self._adguard_turn_off()
+            await self.entity_description.turn_off_fn(self.adguard)()
         except AdGuardHomeError:
             _LOGGER.error("An error occurred while turning off AdGuard Home switch")
             self._available = False
 
-    async def _adguard_turn_off(self) -> None:
-        """Turn off the switch."""
-        raise NotImplementedError()
-
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn on the switch."""
         try:
-            await self._adguard_turn_on()
+            await self.entity_description.turn_on_fn(self.adguard)()
         except AdGuardHomeError:
             _LOGGER.error("An error occurred while turning on AdGuard Home switch")
             self._available = False
 
-    async def _adguard_turn_on(self) -> None:
-        """Turn on the switch."""
-        raise NotImplementedError()
-
-
-class AdGuardHomeProtectionSwitch(AdGuardHomeSwitch):
-    """Defines a AdGuard Home protection switch."""
-
-    def __init__(self, adguard: AdGuardHome, entry: ConfigEntry) -> None:
-        """Initialize AdGuard Home switch."""
-        super().__init__(adguard, entry, "Protection", "mdi:shield-check", "protection")
-
-    async def _adguard_turn_off(self) -> None:
-        """Turn off the switch."""
-        await self.adguard.disable_protection()
-
-    async def _adguard_turn_on(self) -> None:
-        """Turn on the switch."""
-        await self.adguard.enable_protection()
-
     async def _adguard_update(self) -> None:
         """Update AdGuard Home entity."""
-        self._state = await self.adguard.protection_enabled()
+        self._attr_is_on = await self.entity_description.is_on_fn(self.adguard)()
-
-
-class AdGuardHomeParentalSwitch(AdGuardHomeSwitch):
-    """Defines a AdGuard Home parental control switch."""
-
-    def __init__(self, adguard: AdGuardHome, entry: ConfigEntry) -> None:
-        """Initialize AdGuard Home switch."""
-        super().__init__(
-            adguard, entry, "Parental control", "mdi:shield-check", "parental"
-        )
-
-    async def _adguard_turn_off(self) -> None:
-        """Turn off the switch."""
-        await self.adguard.parental.disable()
-
-    async def _adguard_turn_on(self) -> None:
-        """Turn on the switch."""
-        await self.adguard.parental.enable()
-
-    async def _adguard_update(self) -> None:
-        """Update AdGuard Home entity."""
-        self._state = await self.adguard.parental.enabled()
-
-
-class AdGuardHomeSafeSearchSwitch(AdGuardHomeSwitch):
-    """Defines a AdGuard Home safe search switch."""
-
-    def __init__(self, adguard: AdGuardHome, entry: ConfigEntry) -> None:
-        """Initialize AdGuard Home switch."""
-        super().__init__(
-            adguard, entry, "Safe search", "mdi:shield-check", "safesearch"
-        )
-
-    async def _adguard_turn_off(self) -> None:
-        """Turn off the switch."""
-        await self.adguard.safesearch.disable()
-
-    async def _adguard_turn_on(self) -> None:
-        """Turn on the switch."""
-        await self.adguard.safesearch.enable()
-
-    async def _adguard_update(self) -> None:
-        """Update AdGuard Home entity."""
-        self._state = await self.adguard.safesearch.enabled()
-
-
-class AdGuardHomeSafeBrowsingSwitch(AdGuardHomeSwitch):
-    """Defines a AdGuard Home safe search switch."""
-
-    def __init__(self, adguard: AdGuardHome, entry: ConfigEntry) -> None:
-        """Initialize AdGuard Home switch."""
-        super().__init__(
-            adguard, entry, "Safe browsing", "mdi:shield-check", "safebrowsing"
-        )
-
-    async def _adguard_turn_off(self) -> None:
-        """Turn off the switch."""
-        await self.adguard.safebrowsing.disable()
-
-    async def _adguard_turn_on(self) -> None:
-        """Turn on the switch."""
-        await self.adguard.safebrowsing.enable()
-
-    async def _adguard_update(self) -> None:
-        """Update AdGuard Home entity."""
-        self._state = await self.adguard.safebrowsing.enabled()
-
-
-class AdGuardHomeFilteringSwitch(AdGuardHomeSwitch):
-    """Defines a AdGuard Home filtering switch."""
-
-    def __init__(self, adguard: AdGuardHome, entry: ConfigEntry) -> None:
-        """Initialize AdGuard Home switch."""
-        super().__init__(adguard, entry, "Filtering", "mdi:shield-check", "filtering")
-
-    async def _adguard_turn_off(self) -> None:
-        """Turn off the switch."""
-        await self.adguard.filtering.disable()
-
-    async def _adguard_turn_on(self) -> None:
-        """Turn on the switch."""
-        await self.adguard.filtering.enable()
-
-    async def _adguard_update(self) -> None:
-        """Update AdGuard Home entity."""
-        self._state = await self.adguard.filtering.enabled()
-
-
-class AdGuardHomeQueryLogSwitch(AdGuardHomeSwitch):
-    """Defines a AdGuard Home query log switch."""
-
-    def __init__(self, adguard: AdGuardHome, entry: ConfigEntry) -> None:
-        """Initialize AdGuard Home switch."""
-        super().__init__(
-            adguard,
-            entry,
-            "Query log",
-            "mdi:shield-check",
-            "querylog",
-            enabled_default=False,
-        )
-
-    async def _adguard_turn_off(self) -> None:
-        """Turn off the switch."""
-        await self.adguard.querylog.disable()
-
-    async def _adguard_turn_on(self) -> None:
-        """Turn on the switch."""
-        await self.adguard.querylog.enable()
-
-    async def _adguard_update(self) -> None:
-        """Update AdGuard Home entity."""
-        self._state = await self.adguard.querylog.enabled()