Use entity base class for NextDNS entities (#146934)

* Add entity module

* Add NextDnsEntityDescription class

* Remove NextDnsEntityDescription

* Create DeviceInfo in entity module

* Use property

homeassistant/components/nextdns/binary_sensor.py

@@ -13,12 +13,11 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntityDescription,
 )
 from homeassistant.const import EntityCategory
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import NextDnsConfigEntry
-from .coordinator import NextDnsUpdateCoordinator
+from .entity import NextDnsEntity
 
 PARALLEL_UPDATES = 0
 
@@ -61,30 +60,14 @@ async def async_setup_entry(
     )
 
 
-class NextDnsBinarySensor(
-    CoordinatorEntity[NextDnsUpdateCoordinator[ConnectionStatus]], BinarySensorEntity
-):
+class NextDnsBinarySensor(NextDnsEntity, BinarySensorEntity):
     """Define an NextDNS binary sensor."""
 
-    _attr_has_entity_name = True
     entity_description: NextDnsBinarySensorEntityDescription
 
-    def __init__(
-        self,
-        coordinator: NextDnsUpdateCoordinator[ConnectionStatus],
-        description: NextDnsBinarySensorEntityDescription,
-    ) -> None:
-        """Initialize."""
-        super().__init__(coordinator)
-        self._attr_device_info = coordinator.device_info
-        self._attr_unique_id = f"{coordinator.profile_id}_{description.key}"
-        self._attr_is_on = description.state(coordinator.data, coordinator.profile_id)
-        self.entity_description = description
-
-    @callback
-    def _handle_coordinator_update(self) -> None:
-        """Handle updated data from the coordinator."""
-        self._attr_is_on = self.entity_description.state(
+    @property
+    def is_on(self) -> bool:
+        """Return True if the binary sensor is on."""
+        return self.entity_description.state(
             self.coordinator.data, self.coordinator.profile_id
         )
-        self.async_write_ha_state()

homeassistant/components/nextdns/button.py

@@ -4,21 +4,21 @@ from __future__ import annotations
 
 from aiohttp import ClientError
 from aiohttp.client_exceptions import ClientConnectorError
-from nextdns import AnalyticsStatus, ApiError, InvalidApiKeyError
+from nextdns import ApiError, InvalidApiKeyError
 
 from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import NextDnsConfigEntry
 from .const import DOMAIN
-from .coordinator import NextDnsUpdateCoordinator
+from .entity import NextDnsEntity
 
 PARALLEL_UPDATES = 1
 
+
 CLEAR_LOGS_BUTTON = ButtonEntityDescription(
     key="clear_logs",
     translation_key="clear_logs",
@@ -37,24 +37,9 @@ async def async_setup_entry(
     async_add_entities([NextDnsButton(coordinator, CLEAR_LOGS_BUTTON)])
 
 
-class NextDnsButton(
-    CoordinatorEntity[NextDnsUpdateCoordinator[AnalyticsStatus]], ButtonEntity
-):
+class NextDnsButton(NextDnsEntity, ButtonEntity):
     """Define an NextDNS button."""
 
-    _attr_has_entity_name = True
-
-    def __init__(
-        self,
-        coordinator: NextDnsUpdateCoordinator[AnalyticsStatus],
-        description: ButtonEntityDescription,
-    ) -> None:
-        """Initialize."""
-        super().__init__(coordinator)
-        self._attr_device_info = coordinator.device_info
-        self._attr_unique_id = f"{coordinator.profile_id}_{description.key}"
-        self.entity_description = description
-
     async def async_press(self) -> None:
         """Trigger cleaning logs."""
         try:

homeassistant/components/nextdns/coordinator.py

@@ -24,7 +24,6 @@ from tenacity import RetryError
 
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed
-from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 if TYPE_CHECKING:
@@ -53,14 +52,6 @@ class NextDnsUpdateCoordinator(DataUpdateCoordinator[CoordinatorDataT]):
         """Initialize."""
         self.nextdns = nextdns
         self.profile_id = profile_id
-        self.profile_name = nextdns.get_profile_name(profile_id)
-        self.device_info = DeviceInfo(
-            configuration_url=f"https://my.nextdns.io/{profile_id}/setup",
-            entry_type=DeviceEntryType.SERVICE,
-            identifiers={(DOMAIN, str(profile_id))},
-            manufacturer="NextDNS Inc.",
-            name=self.profile_name,
-        )
 
         super().__init__(
             hass,

homeassistant/components/nextdns/entity.py

@@ -0,0 +1,31 @@
+"""Define NextDNS entities."""
+
+from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
+from homeassistant.helpers.entity import EntityDescription
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+
+from .const import DOMAIN
+from .coordinator import CoordinatorDataT, NextDnsUpdateCoordinator
+
+
+class NextDnsEntity(CoordinatorEntity[NextDnsUpdateCoordinator[CoordinatorDataT]]):
+    """Define NextDNS entity."""
+
+    _attr_has_entity_name = True
+
+    def __init__(
+        self,
+        coordinator: NextDnsUpdateCoordinator[CoordinatorDataT],
+        description: EntityDescription,
+    ) -> None:
+        """Initialize."""
+        super().__init__(coordinator)
+        self._attr_device_info = DeviceInfo(
+            configuration_url=f"https://my.nextdns.io/{coordinator.profile_id}/setup",
+            entry_type=DeviceEntryType.SERVICE,
+            identifiers={(DOMAIN, str(coordinator.profile_id))},
+            manufacturer="NextDNS Inc.",
+            name=coordinator.nextdns.get_profile_name(coordinator.profile_id),
+        )
+        self._attr_unique_id = f"{coordinator.profile_id}_{description.key}"
+        self.entity_description = description

homeassistant/components/nextdns/sensor.py

@@ -20,10 +20,9 @@ from homeassistant.components.sensor import (
     SensorStateClass,
 )
 from homeassistant.const import PERCENTAGE, EntityCategory
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.typing import StateType
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import NextDnsConfigEntry
 from .const import (
@@ -33,7 +32,8 @@ from .const import (
     ATTR_PROTOCOLS,
     ATTR_STATUS,
 )
-from .coordinator import CoordinatorDataT, NextDnsUpdateCoordinator
+from .coordinator import CoordinatorDataT
+from .entity import NextDnsEntity
 
 PARALLEL_UPDATES = 0
 
@@ -297,27 +297,12 @@ async def async_setup_entry(
     )
 
 
-class NextDnsSensor(
-    CoordinatorEntity[NextDnsUpdateCoordinator[CoordinatorDataT]], SensorEntity
-):
+class NextDnsSensor(NextDnsEntity, SensorEntity):
     """Define an NextDNS sensor."""
 
-    _attr_has_entity_name = True
+    entity_description: NextDnsSensorEntityDescription
 
-    def __init__(
-        self,
-        coordinator: NextDnsUpdateCoordinator[CoordinatorDataT],
-        description: NextDnsSensorEntityDescription,
-    ) -> None:
-        """Initialize."""
-        super().__init__(coordinator)
-        self._attr_device_info = coordinator.device_info
-        self._attr_unique_id = f"{coordinator.profile_id}_{description.key}"
-        self._attr_native_value = description.value(coordinator.data)
-        self.entity_description: NextDnsSensorEntityDescription = description
-
-    @callback
-    def _handle_coordinator_update(self) -> None:
-        """Handle updated data from the coordinator."""
-        self._attr_native_value = self.entity_description.value(self.coordinator.data)
-        self.async_write_ha_state()
+    @property
+    def native_value(self) -> StateType:
+        """Return the state of the sensor."""
+        return self.entity_description.value(self.coordinator.data)

homeassistant/components/nextdns/switch.py

@@ -15,11 +15,11 @@ from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import NextDnsConfigEntry
 from .const import DOMAIN
 from .coordinator import NextDnsUpdateCoordinator
+from .entity import NextDnsEntity
 
 PARALLEL_UPDATES = 1
 
@@ -536,12 +536,9 @@ async def async_setup_entry(
     )
 
 
-class NextDnsSwitch(
-    CoordinatorEntity[NextDnsUpdateCoordinator[Settings]], SwitchEntity
-):
+class NextDnsSwitch(NextDnsEntity, SwitchEntity):
     """Define an NextDNS switch."""
 
-    _attr_has_entity_name = True
     entity_description: NextDnsSwitchEntityDescription
 
     def __init__(
@@ -550,11 +547,8 @@ class NextDnsSwitch(
         description: NextDnsSwitchEntityDescription,
     ) -> None:
         """Initialize."""
-        super().__init__(coordinator)
-        self._attr_device_info = coordinator.device_info
-        self._attr_unique_id = f"{coordinator.profile_id}_{description.key}"
+        super().__init__(coordinator, description)
         self._attr_is_on = description.state(coordinator.data)
-        self.entity_description = description
 
     @callback
     def _handle_coordinator_update(self) -> None: