some code changes and hotfix for sensors

custom_components/technitiumdns/api.py

@@ -1,7 +1,7 @@
-import aiohttp
 import asyncio
-import async_timeout
 import logging
+import aiohttp
+import async_timeout
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -111,10 +111,10 @@ async def set_ad_blocking(self, enable):
                         return data
             except aiohttp.ClientError as err:
                 _LOGGER.error("Error setting ad blocking: %s", err)
-                raise Exception(f"Error setting ad blocking: {err}")
-            except asyncio.TimeoutError:
+                raise Exception(f"Error setting ad blocking: {err}") from err
+            except asyncio.TimeoutError as e:
                 _LOGGER.error("Timeout error setting ad blocking")
-                raise Exception("Timeout error setting ad blocking")
+                raise Exception("Timeout error setting ad blocking") from e
             except Exception as e:
                 _LOGGER.error("An error occurred: %s", e)
-                raise Exception(f"An error occurred: {e}")
+                raise Exception(f"An error occurred: {e}") from e

custom_components/technitiumdns/config_flow.py

@@ -1,85 +1,86 @@
-"""Config flow for TechnitiumDNS integration."""
-
-from homeassistant import config_entries
-from homeassistant.core import callback
-from homeassistant.helpers import aiohttp_client
-import aiohttp
-import async_timeout
-import voluptuous as vol
-
-from .const import DOMAIN
-from .api import TechnitiumDNSApi
-
-
-@config_entries.HANDLERS.register(DOMAIN)
-class TechnitiumDNSConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
-    """Handle a config flow for TechnitiumDNS."""
-
-    VERSION = 1
-
-    async def async_step_user(self, user_input=None):
-        """Handle the initial step."""
-        errors = {}
-        if user_input is not None:
-            try:
-                # Validate the input by trying to create the API object
-                api = TechnitiumDNSApi(user_input["api_url"], user_input["token"])
-                await api.get_statistics(user_input["stats_duration"])
-
-                return self.async_create_entry(
-                    title=user_input["server_name"], data=user_input
-                )
-            except Exception as e:
-                errors["base"] = "auth"
-
-        data_schema = vol.Schema(
-            {
-                vol.Required("api_url"): str,
-                vol.Required("token"): str,
-                vol.Required("server_name"): str,
-                vol.Required("username"): str,
-                vol.Required("stats_duration"): vol.In(
-                    ["LastHour", "LastDay", "LastWeek", "LastMonth"]
-                ),
-            }
-        )
-
-        return self.async_show_form(
-            step_id="user", data_schema=data_schema, errors=errors
-        )
-
-    async def _test_credentials(self, api_url, token, stats_duration):
-        """Test the provided credentials."""
-        try:
-            session = aiohttp_client.async_get_clientsession(self.hass)
-            with async_timeout.timeout(10):
-                response = await session.get(
-                    f"{api_url}/api/dashboard/stats/get?token={token}&type={stats_duration}&utc=true"
-                )
-                if (
-                    response.status == 200
-                    and (await response.json()).get("status") == "ok"
-                ):
-                    return True
-        except (aiohttp.ClientError, asyncio.TimeoutError):
-            pass
-        return False
-
-    @staticmethod
-    @callback
-    def async_get_options_flow(config_entry):
-        return TechnitiumDNSOptionsFlowHandler(config_entry)
-
-
-class TechnitiumDNSOptionsFlowHandler(config_entries.OptionsFlow):
-    """Handle an options flow for TechnitiumDNS."""
-
-    def __init__(self, config_entry):
-        self.config_entry = config_entry
-
-    async def async_step_init(self, user_input=None):
-        """Manage the options."""
-        if user_input is not None:
-            return self.async_create_entry(title="", data=user_input)
-
-        return self.async_show_form(step_id="init", data_schema=vol.Schema({}))
+"""Config flow for TechnitiumDNS integration."""
+
+
+import asyncio
+from homeassistant import config_entries
+from homeassistant.core import callback
+from homeassistant.helpers import aiohttp_client
+import aiohttp
+import async_timeout
+import contextlib
+import voluptuous as vol
+
+from .const import DOMAIN
+from .api import TechnitiumDNSApi
+
+
+@config_entries.HANDLERS.register(DOMAIN)
+class TechnitiumDNSConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
+    """Handle a config flow for TechnitiumDNS."""
+
+    VERSION = 1
+
+    async def async_step_user(self, user_input=None):
+        """Handle the initial step."""
+        errors = {}
+        if user_input is not None:
+            try:
+                # Validate the input by trying to create the API object
+                api = TechnitiumDNSApi(user_input["api_url"], user_input["token"])
+                await api.get_statistics(user_input["stats_duration"])
+
+                return self.async_create_entry(
+                    title=user_input["server_name"], data=user_input
+                )
+            except Exception as e:
+                errors["base"] = "auth"
+
+        data_schema = vol.Schema(
+            {
+                vol.Required("api_url"): str,
+                vol.Required("token"): str,
+                vol.Required("server_name"): str,
+                vol.Required("username"): str,
+                vol.Required("stats_duration"): vol.In(
+                    ["LastHour", "LastDay", "LastWeek", "LastMonth"]
+                ),
+            }
+        )
+
+        return self.async_show_form(
+            step_id="user", data_schema=data_schema, errors=errors
+        )
+
+    async def _test_credentials(self, api_url, token, stats_duration):
+        """Test the provided credentials."""
+        with contextlib.suppress(aiohttp.ClientError, asyncio.TimeoutError):
+            session = aiohttp_client.async_get_clientsession(self.hass)
+            with async_timeout.timeout(10):
+                response = await session.get(
+                    f"{api_url}/api/dashboard/stats/get?token={token}&type={stats_duration}&utc=true"
+                )
+                if (
+                    response.status == 200
+                    and (await response.json()).get("status") == "ok"
+                ):
+                    return True
+        return False
+
+    @staticmethod
+    @callback
+    def async_get_options_flow(config_entry):
+        return TechnitiumDNSOptionsFlowHandler(config_entry)
+
+
+class TechnitiumDNSOptionsFlowHandler(config_entries.OptionsFlow):
+    """Handle an options flow for TechnitiumDNS."""
+
+    def __init__(self, config_entry):
+        self.config_entry = config_entry
+
+    async def async_step_init(self, user_input=None):
+        """Manage the options."""
+        if user_input is not None:
+            return self.async_create_entry(title="", data=user_input)
+
+        return self.async_show_form(step_id="init", data_schema=vol.Schema({}))

custom_components/technitiumdns/const.py

@@ -1,39 +1,39 @@
-"""Constants for the TechnitiumDNS integration."""
-
-DOMAIN = "technitiumdns"
-
-DURATION_OPTIONS = ["LastHour", "LastDay", "LastWeek", "LastMonth"]
-
-AD_BLOCKING_SWITCH = "Enable Ad Blocking"
-
-AD_BLOCKING_DURATION_OPTIONS = {
-    5: "Disable Ad Blocking for 5 Minutes",
-    10: "Disable Ad Blocking for 10 Minutes",
-    30: "Disable Ad Blocking for 30 Minutes",
-    60: "Disable Ad Blocking for 60 Minutes",
-    1440: "Disable Ad Blocking for 1 Day",
-}
-
-SENSOR_TYPES = {
-    "queries": {"name": "Total Queries", "state_class": "measurement", "device_class": "count"},
-    "blocked_queries": {"name": "Blocked Queries", "state_class": "measurement", "device_class": "count"},
-    "clients": {"name": "Clients", "state_class": "measurement", "device_class": "count"},
-    "update_available": {"name": "Update Available", "state_class": None, "device_class": "update"},
-    "no_error": {"name": "Total No Error", "state_class": "measurement", "device_class": "count"},
-    "server_failure": {"name": "Total Server Failure", "state_class": "measurement", "device_class": "count"},
-    "nx_domain": {"name": "Total NX Domain", "state_class": "measurement", "device_class": "count"},
-    "refused": {"name": "Total Refused", "state_class": "measurement", "device_class": "count"},
-    "authoritative": {"name": "Total Authoritative", "state_class": "measurement", "device_class": "count"},
-    "recursive": {"name": "Total Recursive", "state_class": "measurement", "device_class": "count"},
-    "cached": {"name": "Total Cached", "state_class": "measurement", "device_class": "count"},
-    "dropped": {"name": "Total Dropped", "state_class": "measurement", "device_class": "count"},
-    "zones": {"name": "Zones", "state_class": "measurement", "device_class": "count"},
-    "cached_entries": {"name": "Cached Entries", "state_class": "measurement", "device_class": "count"},
-    "allowed_zones": {"name": "Allowed Zones", "state_class": "measurement", "device_class": "count"},
-    "blocked_zones": {"name": "Blocked Zones", "state_class": "measurement", "device_class": "count"},
-    "allow_list_zones": {"name": "Allow List Zones", "state_class": "measurement", "device_class": "count"},
-    "block_list_zones": {"name": "Block List Zones", "state_class": "measurement", "device_class": "count"},
-    "top_clients": {"name": "Top Clients", "state_class": None, "device_class": None},
-    "top_domains": {"name": "Top Domains", "state_class": None, "device_class": None},
-    "top_blocked_domains": {"name": "Top Blocked Domains", "state_class": None, "device_class": None},
-}
+"""Constants for the TechnitiumDNS integration."""
+
+DOMAIN = "technitiumdns"
+
+DURATION_OPTIONS = ["LastHour", "LastDay", "LastWeek", "LastMonth"]
+
+AD_BLOCKING_SWITCH = "Enable Ad Blocking"
+
+AD_BLOCKING_DURATION_OPTIONS = {
+    5: "Disable Ad Blocking for 5 Minutes",
+    10: "Disable Ad Blocking for 10 Minutes",
+    30: "Disable Ad Blocking for 30 Minutes",
+    60: "Disable Ad Blocking for 60 Minutes",
+    1440: "Disable Ad Blocking for 1 Day",
+}
+
+SENSOR_TYPES = {
+    "queries": {"name": "Total Queries", "state_class": "measurement", "device_class": "count"},
+    "blocked_queries": {"name": "Blocked Queries", "state_class": "measurement", "device_class": "count"},
+    "clients": {"name": "Clients", "state_class": "measurement", "device_class": "count"},
+    "update_available": {"name": "Update Available", "state_class": None, "device_class": "update"},
+    "no_error": {"name": "Total No Error", "state_class": "measurement", "device_class": "count"},
+    "server_failure": {"name": "Total Server Failure", "state_class": "measurement", "device_class": "count"},
+    "nx_domain": {"name": "Total NX Domain", "state_class": "measurement", "device_class": "count"},
+    "refused": {"name": "Total Refused", "state_class": "measurement", "device_class": "count"},
+    "authoritative": {"name": "Total Authoritative", "state_class": "measurement", "device_class": "count"},
+    "recursive": {"name": "Total Recursive", "state_class": "measurement", "device_class": "count"},
+    "cached": {"name": "Total Cached", "state_class": "measurement", "device_class": "count"},
+    "dropped": {"name": "Total Dropped", "state_class": "measurement", "device_class": "count"},
+    "zones": {"name": "Zones", "state_class": "measurement", "device_class": "count"},
+    "cached_entries": {"name": "Cached Entries", "state_class": "measurement", "device_class": "count"},
+    "allowed_zones": {"name": "Allowed Zones", "state_class": "measurement", "device_class": "count"},
+    "blocked_zones": {"name": "Blocked Zones", "state_class": "measurement", "device_class": "count"},
+    "allow_list_zones": {"name": "Allow List Zones", "state_class": "measurement", "device_class": "count"},
+    "block_list_zones": {"name": "Block List Zones", "state_class": "measurement", "device_class": "count"},
+    "top_clients": {"name": "Top Clients", "state_class": None, "device_class": None},
+    "top_domains": {"name": "Top Domains", "state_class": None, "device_class": None},
+    "top_blocked_domains": {"name": "Top Blocked Domains", "state_class": None, "device_class": None},
+}

custom_components/technitiumdns/sensor.py

@@ -141,8 +141,14 @@ def state(self):
     @property
     def extra_state_attributes(self):
         """Return additional attributes in a table-friendly format based on sensor type."""
-        attributes = {}
-
+        attributes = {
+            "queries": self.coordinator.data.get("queries"),
+            "blocked_queries": self.coordinator.data.get("blocked_queries"),
+            "clients": self.coordinator.data.get("clients"),
+            "update_available": self.coordinator.data.get("update_available"),
+        }
+
+        # Add structured table data based on the sensor type
         if self._sensor_type == 'top_clients':
             attributes["top_clients_table"] = [
                 {"Client": client.get('name', 'Unknown'), "Hits": client.get('hits', 0)}