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

added device registration

custom_components/technitiumdns/__init__.py

@@ -4,6 +4,7 @@
 
 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
@@ -21,7 +22,16 @@ 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)},
+        name=entry.data["server_name"],
+        manufacturer="Technitium",
+        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"]
     )

custom_components/technitiumdns/button.py

@@ -11,16 +11,21 @@
 
 async def async_setup_entry(hass: HomeAssistant, 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
     ]
@@ -34,17 +39,17 @@ class TechnitiumDNSButton(ButtonEntity):
 
     def __init__(
         self,
-        hass: HomeAssistant,
         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."""
@@ -55,3 +60,14 @@ async def async_press(self) -> None:
             )
         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 {
+            "identifiers": {(DOMAIN, self._entry_id)},
+            "name": self._attr_name,
+            "manufacturer": "Technitium",
+            "model": "DNS Server",
+            "entry_type": "service",
+        }

custom_components/technitiumdns/sensor.py

@@ -7,6 +7,7 @@
     UpdateFailed,
     CoordinatorEntity,
 )
+from homeassistant.helpers import device_registry as dr
 
 from .const import DOMAIN, SENSOR_TYPES
 from .api import TechnitiumDNSApi
@@ -26,9 +27,14 @@ async def async_setup_entry(hass, entry, async_add_entities):
     coordinator = TechnitiumDNSCoordinator(hass, api, stats_duration)
     await coordinator.async_config_entry_first_refresh()
 
+    device_registry = dr.async_get(hass)
+    device = device_registry.async_get_device(identifiers={(DOMAIN, entry.entry_id)})
+
     sensors = []
     for sensor_type in SENSOR_TYPES:
-        sensors.append(TechnitiumDNSSensor(coordinator, sensor_type, server_name))
+        sensors.append(
+            TechnitiumDNSSensor(coordinator, sensor_type, server_name, device.id)
+        )
 
     async_add_entities(sensors, True)
 
@@ -141,11 +147,12 @@ async def _async_update_data(self):
 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, device_id):
         """Initialize the sensor."""
         super().__init__(coordinator)
         self._sensor_type = sensor_type
         self._server_name = server_name
+        self._device_id = device_id
         self._name = (
             f"Technitiumdns_{SENSOR_TYPES[sensor_type]['name']} ({server_name})"
         )
@@ -188,3 +195,13 @@ def available(self):
     def should_poll(self):
         """No polling needed."""
         return False
+
+    @property
+    def device_info(self):
+        """Return the device info."""
+        return {
+            "identifiers": {(DOMAIN, self._device_id)},
+            "name": self._server_name,
+            "manufacturer": "Technitium",
+            "model": "DNS Server",
+        }

custom_components/technitiumdns/switch.py

@@ -11,11 +11,14 @@
 
 async def async_setup_entry(hass: HomeAssistant, 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)
@@ -25,13 +28,13 @@ class TechnitiumDNSSwitch(SwitchEntity):
     """Representation of a TechnitiumDNS switch."""
 
     def __init__(
-        self, hass: HomeAssistant, api: TechnitiumDNSApi, name: str, server_name: str
+        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):
@@ -78,3 +81,14 @@ 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 {
+            "identifiers": {(DOMAIN, self._entry_id)},
+            "name": self._attr_name,
+            "manufacturer": "Technitium",
+            "model": "DNS Server",
+            "entry_type": "service",
+        }