Skip to content

Commit

Permalink
possible hotfix for registering the device twice
Browse files Browse the repository at this point in the history
  • Loading branch information
@Amateur-God
Amateur-God committed Jun 22, 2024
1 parent 2a560c2 commit 63d8342
Show file tree
Hide file tree
Showing 4 changed files with 51 additions and 134 deletions.
4 changes: 1 addition & 3 deletions custom_components/technitiumdns/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@

_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, {})
Expand All @@ -26,8 +25,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
device_registry.async_get_or_create(
config_entry_id=entry.entry_id,
identifiers={(DOMAIN, entry.entry_id)},
name=entry.data["server_name"],
manufacturer="Technitium",
name=entry.data["server_name"],
model="DNS Server",
)

Expand All @@ -37,7 +36,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
)
return True


async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Unload a config entry."""
await hass.config_entries.async_unload_platforms(
Expand Down
39 changes: 11 additions & 28 deletions custom_components/technitiumdns/button.py
View file
Original file line number Diff line number Diff line change
@@ -1,15 +1,13 @@
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."""
config_entry = hass.data[DOMAIN][entry.entry_id]
api = config_entry["api"]
Expand All @@ -21,30 +19,18 @@ async def async_setup_entry(hass: HomeAssistant, entry, async_add_entities):
# Define the buttons using the sorted durations
buttons = [
TechnitiumDNSButton(
api,
AD_BLOCKING_DURATION_OPTIONS[duration],
duration,
server_name,
entry.entry_id,
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,
api: TechnitiumDNSApi,
name: str,
duration: int,
server_name: str,
entry_id: str,
):
def __init__(self, api: TechnitiumDNSApi, name: str, duration: int, server_name: str, entry_id: str):
"""Initialize the button."""
self._api = api
self._attr_name = f"{name} ({server_name})"
Expand All @@ -55,19 +41,16 @@ 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 {
"identifiers": {(DOMAIN, self._entry_id)},
"name": self._attr_name,
"manufacturer": "Technitium",
"model": "DNS Server",
"entry_type": "service",
}
return DeviceInfo(
identifiers={(DOMAIN, self._entry_id)},
name=self._attr_name,
manufacturer="Technitium",
model="DNS Server",
)
110 changes: 28 additions & 82 deletions custom_components/technitiumdns/sensor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
UpdateFailed,
CoordinatorEntity,
)
from homeassistant.helpers import device_registry as dr
from homeassistant.helpers.entity import DeviceInfo

from .const import DOMAIN, SENSOR_TYPES
from .api import TechnitiumDNSApi
Expand All @@ -16,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]
Expand All @@ -27,18 +26,12 @@ 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, device.id)
)
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."""

Expand All @@ -52,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"),
Expand All @@ -113,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)
Expand All @@ -143,19 +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, device_id):
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._device_id = device_id
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):
Expand All @@ -170,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)):
Expand All @@ -198,10 +144,10 @@ def should_poll(self):

@property
def device_info(self):
"""Return the device info."""
return {
"identifiers": {(DOMAIN, self._device_id)},
"name": self._server_name,
"manufacturer": "Technitium",
"model": "DNS Server",
}
"""Return device information for this entity."""
return DeviceInfo(
identifiers={(DOMAIN, self._entry_id)},
name=self._server_name,
manufacturer="Technitium",
model="DNS Server",
)
32 changes: 11 additions & 21 deletions custom_components/technitiumdns/switch.py
View file
Original file line number Diff line number Diff line change
@@ -1,35 +1,28 @@
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."""
config_entry = hass.data[DOMAIN][entry.entry_id]
api = config_entry["api"]
server_name = config_entry["server_name"]

# Define the switch
switches = [
TechnitiumDNSSwitch(api, AD_BLOCKING_SWITCH, server_name, entry.entry_id)
]
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, api: TechnitiumDNSApi, name: str, server_name: str, entry_id: str
):
def __init__(self, api: TechnitiumDNSApi, name: str, server_name: str, entry_id: str):
"""Initialize the switch."""
self._api = api
self._attr_name = f"{name} ({server_name})"
Expand All @@ -55,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}")
Expand Down Expand Up @@ -85,10 +76,9 @@ async def async_turn_off(self, **kwargs):
@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",
}
return DeviceInfo(
identifiers={(DOMAIN, self._entry_id)},
name=self._attr_name,
manufacturer="Technitium",
model="DNS Server",
)

0 comments on commit 63d8342

Please sign in to comment.