Merge pull request #27 from Amateur-God/V2.2.0-Beta

V2.2.0 beta

custom_components/technitiumdns/__init__.py

@@ -4,13 +4,13 @@
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr
 
 from .const import DOMAIN
 from .api import TechnitiumDNSApi
 
 _LOGGER = logging.getLogger(__name__)
 
-
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up TechnitiumDNS from a config entry."""
     hass.data.setdefault(DOMAIN, {})
@@ -21,13 +21,21 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         "stats_duration": entry.data["stats_duration"],
     }
 
-    # Forward the setup to the sensor and switch platforms
+    device_registry = dr.async_get(hass)
+    device_registry.async_get_or_create(
+        config_entry_id=entry.entry_id,
+        identifiers={(DOMAIN, entry.entry_id)},
+        manufacturer="Technitium",
+        name=entry.data["server_name"],
+        model="DNS Server",
+    )
+
+    # Forward the setup to the sensor, button, and switch platforms
     await hass.config_entries.async_forward_entry_setups(
         entry, ["sensor", "button", "switch"]
     )
     return True
 
-
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     await hass.config_entries.async_unload_platforms(

custom_components/technitiumdns/button.py

@@ -1,57 +1,56 @@
 import logging
 from homeassistant.components.button import ButtonEntity
-from homeassistant.core import HomeAssistant
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.helpers.entity import DeviceInfo
 
 from .const import DOMAIN, AD_BLOCKING_DURATION_OPTIONS
 from .api import TechnitiumDNSApi
 
 _LOGGER = logging.getLogger(__name__)
 
-
-async def async_setup_entry(hass: HomeAssistant, entry, async_add_entities):
+async def async_setup_entry(hass, entry, async_add_entities):
     """Set up TechnitiumDNS button entities based on a config entry."""
-    api = TechnitiumDNSApi(entry.data["api_url"], entry.data["token"])
-    server_name = entry.data["server_name"]
+    config_entry = hass.data[DOMAIN][entry.entry_id]
+    api = config_entry["api"]
+    server_name = config_entry["server_name"]
 
     # Ensure durations are sorted as integers
     sorted_durations = sorted(AD_BLOCKING_DURATION_OPTIONS.keys())
 
     # Define the buttons using the sorted durations
     buttons = [
         TechnitiumDNSButton(
-            hass, api, AD_BLOCKING_DURATION_OPTIONS[duration], duration, server_name
+            api, AD_BLOCKING_DURATION_OPTIONS[duration], duration, server_name, entry.entry_id
         )
         for duration in sorted_durations
     ]
 
     # Add entities
     async_add_entities(buttons)
 
-
 class TechnitiumDNSButton(ButtonEntity):
     """Representation of a TechnitiumDNS button."""
 
-    def __init__(
-        self,
-        hass: HomeAssistant,
-        api: TechnitiumDNSApi,
-        name: str,
-        duration: int,
-        server_name: str,
-    ):
+    def __init__(self, api: TechnitiumDNSApi, name: str, duration: int, server_name: str, entry_id: str):
         """Initialize the button."""
-        self._hass = hass
         self._api = api
         self._attr_name = f"{name} ({server_name})"
         self._duration = duration
+        self._entry_id = entry_id
 
     async def async_press(self) -> None:
         """Handle the button press."""
         try:
             await self._api.temporary_disable_blocking(self._duration)
-            _LOGGER.info(
-                f"Ad blocking disabled for {self._duration} minutes on {self._attr_name}"
-            )
+            _LOGGER.info(f"Ad blocking disabled for {self._duration} minutes on {self._attr_name}")
         except Exception as e:
             _LOGGER.error(f"Failed to disable ad blocking: {e}")
+
+    @property
+    def device_info(self):
+        """Return device information for this entity."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, self._entry_id)},
+            name=self._attr_name,
+            manufacturer="Technitium",
+            model="DNS Server",
+        )

custom_components/technitiumdns/manifest.json

@@ -11,5 +11,5 @@
   "iot_class": "local_polling",
   "issue_tracker": "https://github.com/Amateur-God/home-assistant-Technitiumdns/issues",
   "requirements": ["aiohttp==3.9.5"],
-  "version": "2.1.2"
+  "version": "2.2.0"
 }

custom_components/technitiumdns/sensor.py

@@ -7,6 +7,7 @@
     UpdateFailed,
     CoordinatorEntity,
 )
+from homeassistant.helpers.entity import DeviceInfo
 
 from .const import DOMAIN, SENSOR_TYPES
 from .api import TechnitiumDNSApi
@@ -15,7 +16,6 @@
 
 SCAN_INTERVAL = timedelta(minutes=1)
 
-
 async def async_setup_entry(hass, entry, async_add_entities):
     """Set up the TechnitiumDNS sensor based on a config entry."""
     config_entry = hass.data[DOMAIN][entry.entry_id]
@@ -28,11 +28,10 @@ async def async_setup_entry(hass, entry, async_add_entities):
 
     sensors = []
     for sensor_type in SENSOR_TYPES:
-        sensors.append(TechnitiumDNSSensor(coordinator, sensor_type, server_name))
+        sensors.append(TechnitiumDNSSensor(coordinator, sensor_type, server_name, entry.entry_id))
 
     async_add_entities(sensors, True)
 
-
 class TechnitiumDNSCoordinator(DataUpdateCoordinator):
     """Class to manage fetching TechnitiumDNS data."""
 
@@ -46,52 +45,25 @@ async def _async_update_data(self):
         """Update data via library."""
         try:
             _LOGGER.debug("Fetching data from TechnitiumDNS API")
-            Technitiumdns_statistics = await self.api.get_statistics(
-                self.stats_duration
-            )
-            Technitiumdns_top_clients = await self.api.get_top_clients(
-                self.stats_duration
-            )
-            Technitiumdns_top_domains = await self.api.get_top_domains(
-                self.stats_duration
-            )
-            Technitiumdns_top_blocked_domains = await self.api.get_top_blocked_domains(
-                self.stats_duration
-            )
+            Technitiumdns_statistics = await self.api.get_statistics(self.stats_duration)
+            Technitiumdns_top_clients = await self.api.get_top_clients(self.stats_duration)
+            Technitiumdns_top_domains = await self.api.get_top_domains(self.stats_duration)
+            Technitiumdns_top_blocked_domains = await self.api.get_top_blocked_domains(self.stats_duration)
             Technitiumdns_update_info = await self.api.check_update()
 
             # Add more logging to debug empty response issue
-            _LOGGER.debug(
-                "Technitiumdns_statistics response content: %s",
-                Technitiumdns_statistics,
-            )
-            _LOGGER.debug(
-                "Technitiumdns_top_clients response content: %s",
-                Technitiumdns_top_clients,
-            )
-            _LOGGER.debug(
-                "Technitiumdns_top_domains response content: %s",
-                Technitiumdns_top_domains,
-            )
-            _LOGGER.debug(
-                "Technitiumdns_top_blocked_domains response content: %s",
-                Technitiumdns_top_blocked_domains,
-            )
-            _LOGGER.debug(
-                "Technitiumdns_update_info response content: %s",
-                Technitiumdns_update_info,
-            )
-
-            Technitiumdns_stats = Technitiumdns_statistics.get("response", {}).get(
-                "stats", {}
-            )
+            _LOGGER.debug("Technitiumdns_statistics response content: %s", Technitiumdns_statistics)
+            _LOGGER.debug("Technitiumdns_top_clients response content: %s", Technitiumdns_top_clients)
+            _LOGGER.debug("Technitiumdns_top_domains response content: %s", Technitiumdns_top_domains)
+            _LOGGER.debug("Technitiumdns_top_blocked_domains response content: %s", Technitiumdns_top_blocked_domains)
+            _LOGGER.debug("Technitiumdns_update_info response content: %s", Technitiumdns_update_info)
+
+            Technitiumdns_stats = Technitiumdns_statistics.get("response", {}).get("stats", {})
             data = {
                 "queries": Technitiumdns_stats.get("totalQueries"),
                 "blocked_queries": Technitiumdns_stats.get("totalBlocked"),
                 "clients": Technitiumdns_stats.get("totalClients"),
-                "update_available": Technitiumdns_update_info.get("response", {}).get(
-                    "updateAvailable"
-                ),
+                "update_available": Technitiumdns_update_info.get("response", {}).get("updateAvailable"),
                 "no_error": Technitiumdns_stats.get("totalNoError"),
                 "server_failure": Technitiumdns_stats.get("totalServerFailure"),
                 "nx_domain": Technitiumdns_stats.get("totalNxDomain"),
@@ -107,28 +79,13 @@ async def _async_update_data(self):
                 "allow_list_zones": Technitiumdns_stats.get("allowListZones"),
                 "block_list_zones": Technitiumdns_stats.get("blockListZones"),
                 "top_clients": "\n".join(
-                    [
-                        f"{client['name']} ({client['hits']})"
-                        for client in Technitiumdns_top_clients.get("response", {}).get(
-                            "topClients", []
-                        )[:5]
-                    ]
+                    [f"{client['name']} ({client['hits']})" for client in Technitiumdns_top_clients.get("response", {}).get("topClients", [])[:5]]
                 ),
                 "top_domains": "\n".join(
-                    [
-                        f"{domain['name']} ({domain['hits']})"
-                        for domain in Technitiumdns_top_domains.get("response", {}).get(
-                            "topDomains", []
-                        )[:5]
-                    ]
+                    [f"{domain['name']} ({domain['hits']})" for domain in Technitiumdns_top_domains.get("response", {}).get("topDomains", [])[:5]]
                 ),
                 "top_blocked_domains": "\n".join(
-                    [
-                        f"{domain['name']} ({domain['hits']})"
-                        for domain in Technitiumdns_top_blocked_domains.get(
-                            "response", {}
-                        ).get("topBlockedDomains", [])[:5]
-                    ]
+                    [f"{domain['name']} ({domain['hits']})" for domain in Technitiumdns_top_blocked_domains.get("response", {}).get("topBlockedDomains", [])[:5]]
                 ),
             }
             _LOGGER.debug("Data combined: %s", data)
@@ -137,18 +94,16 @@ async def _async_update_data(self):
             _LOGGER.error("Error fetching data: %s", err)
             raise UpdateFailed(f"Error fetching data: {err}")
 
-
 class TechnitiumDNSSensor(CoordinatorEntity, SensorEntity):
     """Representation of a TechnitiumDNS sensor."""
 
-    def __init__(self, coordinator, sensor_type, server_name):
+    def __init__(self, coordinator, sensor_type, server_name, entry_id):
         """Initialize the sensor."""
         super().__init__(coordinator)
         self._sensor_type = sensor_type
         self._server_name = server_name
-        self._name = (
-            f"Technitiumdns_{SENSOR_TYPES[sensor_type]['name']} ({server_name})"
-        )
+        self._entry_id = entry_id
+        self._name = f"Technitiumdns_{SENSOR_TYPES[sensor_type]['name']} ({server_name})"
 
     @property
     def name(self):
@@ -163,9 +118,7 @@ def state(self):
 
         # Ensure the state value is within the allowable length
         if isinstance(state_value, str) and len(state_value) > 255:
-            _LOGGER.error(
-                "State value for %s exceeds 255 characters", self._sensor_type
-            )
+            _LOGGER.error("State value for %s exceeds 255 characters", self._sensor_type)
             return state_value[:255]
 
         if isinstance(state_value, (list, dict)):
@@ -188,3 +141,13 @@ def available(self):
     def should_poll(self):
         """No polling needed."""
         return False
+
+    @property
+    def device_info(self):
+        """Return device information for this entity."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, self._entry_id)},
+            name=self._server_name,
+            manufacturer="Technitium",
+            model="DNS Server",
+        )

custom_components/technitiumdns/switch.py

@@ -1,37 +1,33 @@
 import logging
 from homeassistant.components.switch import SwitchEntity
-from homeassistant.core import HomeAssistant
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.helpers.entity import DeviceInfo
 
 from .const import DOMAIN, AD_BLOCKING_SWITCH
 from .api import TechnitiumDNSApi
 
 _LOGGER = logging.getLogger(__name__)
 
-
-async def async_setup_entry(hass: HomeAssistant, entry, async_add_entities):
+async def async_setup_entry(hass, entry, async_add_entities):
     """Set up TechnitiumDNS switch entities based on a config entry."""
-    api = TechnitiumDNSApi(entry.data["api_url"], entry.data["token"])
-    server_name = entry.data["server_name"]
+    config_entry = hass.data[DOMAIN][entry.entry_id]
+    api = config_entry["api"]
+    server_name = config_entry["server_name"]
 
     # Define the switch
-    switches = [TechnitiumDNSSwitch(hass, api, AD_BLOCKING_SWITCH, server_name)]
+    switches = [TechnitiumDNSSwitch(api, AD_BLOCKING_SWITCH, server_name, entry.entry_id)]
 
     # Add entities
     async_add_entities(switches)
 
-
 class TechnitiumDNSSwitch(SwitchEntity):
     """Representation of a TechnitiumDNS switch."""
 
-    def __init__(
-        self, hass: HomeAssistant, api: TechnitiumDNSApi, name: str, server_name: str
-    ):
+    def __init__(self, api: TechnitiumDNSApi, name: str, server_name: str, entry_id: str):
         """Initialize the switch."""
-        self._hass = hass
         self._api = api
         self._attr_name = f"{name} ({server_name})"
         self._is_on = False
+        self._entry_id = entry_id
 
     @property
     def name(self):
@@ -52,9 +48,7 @@ async def _fetch_state(self):
         try:
             response = await self._api.get_dns_settings()
             self._is_on = response["response"].get("enableBlocking", False)
-            _LOGGER.info(
-                f"Fetched ad blocking state: {self._is_on} for {self._attr_name}"
-            )
+            _LOGGER.info(f"Fetched ad blocking state: {self._is_on} for {self._attr_name}")
             self.async_write_ha_state()
         except Exception as e:
             _LOGGER.error(f"Failed to fetch ad blocking state: {e}")
@@ -78,3 +72,13 @@ async def async_turn_off(self, **kwargs):
             self.async_write_ha_state()
         except Exception as e:
             _LOGGER.error(f"Failed to disable ad blocking: {e}")
+
+    @property
+    def device_info(self):
+        """Return device information for this entity."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, self._entry_id)},
+            name=self._attr_name,
+            manufacturer="Technitium",
+            model="DNS Server",
+        )