Use runtime_data and validate connection at setup for dnsip

homeassistant/components/dnsip/__init__.py

@@ -1,26 +1,119 @@
 """The DNS IP integration."""
 
+import asyncio
+from dataclasses import dataclass
+import logging
+
+import aiodns
+from aiodns.error import DNSError
+from pycares import AresError
+
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PORT
-from homeassistant.core import _LOGGER, HomeAssistant
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryNotReady
+
+from .const import (
+    CONF_HOSTNAME,
+    CONF_IPV4,
+    CONF_IPV6,
+    CONF_PORT_IPV6,
+    CONF_RESOLVER,
+    CONF_RESOLVER_IPV6,
+    DEFAULT_PORT,
+    PLATFORMS,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class DnsIPRuntimeData:
+    """Runtime data for DNS IP integration."""
+
+    resolver_ipv4: aiodns.DNSResolver | None
+    resolver_ipv6: aiodns.DNSResolver | None
 
-from .const import CONF_PORT_IPV6, DEFAULT_PORT, PLATFORMS
 
+type DnsIPConfigEntry = ConfigEntry[DnsIPRuntimeData]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: DnsIPConfigEntry) -> bool:
     """Set up DNS IP from a config entry."""
 
+    hostname = entry.data[CONF_HOSTNAME]
+    resolver_ipv4: aiodns.DNSResolver | None = None
+    resolver_ipv6: aiodns.DNSResolver | None = None
+    queries: list = []
+
+    if entry.data[CONF_IPV4]:
+        resolver_ipv4 = aiodns.DNSResolver(
+            nameservers=[entry.options[CONF_RESOLVER]],
+            tcp_port=entry.options[CONF_PORT],
+            udp_port=entry.options[CONF_PORT],
+        )
+        queries.append(resolver_ipv4.query(hostname, "A"))
+
+    if entry.data[CONF_IPV6]:
+        resolver_ipv6 = aiodns.DNSResolver(
+            nameservers=[entry.options[CONF_RESOLVER_IPV6]],
+            tcp_port=entry.options[CONF_PORT_IPV6],
+            udp_port=entry.options[CONF_PORT_IPV6],
+        )
+        queries.append(resolver_ipv6.query(hostname, "AAAA"))
+
+    async def _close_resolvers() -> None:
+        if resolver_ipv4 is not None:
+            await resolver_ipv4.close()
+        if resolver_ipv6 is not None:
+            await resolver_ipv6.close()
+
+    try:
+        async with asyncio.timeout(10):
+            results = await asyncio.gather(*queries, return_exceptions=True)
+    except TimeoutError as err:
+        await _close_resolvers()
+        raise ConfigEntryNotReady(
+            f"DNS lookup timed out for {hostname}: {err}"
+        ) from err
+
+    errors = [
+        result
+        for result in results
+        if isinstance(
+            result, (TimeoutError, DNSError, AresError, asyncio.CancelledError)
+        )
+    ]
+    if errors and len(errors) == len(results):
+        await _close_resolvers()
+        raise ConfigEntryNotReady(
+            f"DNS lookup failed for {hostname}: {errors[0]}"
+        ) from errors[0]
+
+    entry.runtime_data = DnsIPRuntimeData(
+        resolver_ipv4=resolver_ipv4,
+        resolver_ipv6=resolver_ipv6,
+    )
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: DnsIPConfigEntry) -> bool:
     """Unload DNS IP config entry."""
 
-    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    if unload_ok:
+        if entry.runtime_data.resolver_ipv4 is not None:
+            await entry.runtime_data.resolver_ipv4.close()
+        if entry.runtime_data.resolver_ipv6 is not None:
+            await entry.runtime_data.resolver_ipv6.close()
+    return unload_ok
 
 
-async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_migrate_entry(
+    hass: HomeAssistant, config_entry: DnsIPConfigEntry
+) -> bool:
     """Migrate old entry to a newer version."""
 
     if config_entry.version > 1:

homeassistant/components/dnsip/sensor.py

@@ -4,18 +4,18 @@
 from datetime import timedelta
 from ipaddress import IPv4Address, IPv6Address
 import logging
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
 
 import aiodns
 from aiodns.error import DNSError
 
 from homeassistant.components.sensor import SensorEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME, CONF_PORT
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
+from . import DnsIPConfigEntry
 from .const import (
     CONF_HOSTNAME,
     CONF_IPV4,
@@ -46,24 +46,37 @@ def sort_ips(ips: list, querytype: Literal["A", "AAAA"]) -> list:
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: DnsIPConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the dnsip sensor entry."""
 
     hostname = entry.data[CONF_HOSTNAME]
     name = entry.data[CONF_NAME]
 
-    nameserver_ipv4 = entry.options[CONF_RESOLVER]
-    nameserver_ipv6 = entry.options[CONF_RESOLVER_IPV6]
-    port_ipv4 = entry.options[CONF_PORT]
-    port_ipv6 = entry.options[CONF_PORT_IPV6]
-
     entities = []
     if entry.data[CONF_IPV4]:
-        entities.append(WanIpSensor(name, hostname, nameserver_ipv4, False, port_ipv4))
+        entities.append(
+            WanIpSensor(
+                entry,
+                name,
+                hostname,
+                entry.options[CONF_RESOLVER],
+                False,
+                entry.options[CONF_PORT],
+            )
+        )
     if entry.data[CONF_IPV6]:
-        entities.append(WanIpSensor(name, hostname, nameserver_ipv6, True, port_ipv6))
+        entities.append(
+            WanIpSensor(
+                entry,
+                name,
+                hostname,
+                entry.options[CONF_RESOLVER_IPV6],
+                True,
+                entry.options[CONF_PORT_IPV6],
+            )
+        )
 
     async_add_entities(entities, update_before_add=True)
 
@@ -75,17 +88,18 @@ class WanIpSensor(SensorEntity):
     _attr_translation_key = "dnsip"
     _unrecorded_attributes = frozenset({"resolver", "querytype", "ip_addresses"})
 
-    resolver: aiodns.DNSResolver
-
     def __init__(
         self,
+        entry: DnsIPConfigEntry,
         name: str,
         hostname: str,
         nameserver: str,
         ipv6: bool,
         port: int,
     ) -> None:
         """Initialize the DNS IP sensor."""
+        self.entry = entry
+        self.ipv6 = ipv6
         self._attr_name = "IPv6" if ipv6 else None
         self._attr_unique_id = f"{hostname}_{ipv6}"
         self.hostname = hostname
@@ -104,28 +118,43 @@ def __init__(
             model=aiodns.__version__,
             name=name,
         )
-        self.create_dns_resolver()
+
+    @property
+    def _resolver(self) -> aiodns.DNSResolver:
+        """Return the active DNS resolver from runtime data."""
+        resolver = (
+            self.entry.runtime_data.resolver_ipv6
+            if self.ipv6
+            else self.entry.runtime_data.resolver_ipv4
+        )
+        if TYPE_CHECKING:
+            assert resolver is not None
+        return resolver
 
     def create_dns_resolver(self) -> None:
-        """Create the DNS resolver."""
-        self.resolver = aiodns.DNSResolver(
+        """Create a new DNS resolver and store it on runtime data."""
+        new_resolver = aiodns.DNSResolver(
             nameservers=[self.nameserver], tcp_port=self.port, udp_port=self.port
         )
+        if self.ipv6:
+            self.entry.runtime_data.resolver_ipv6 = new_resolver
+        else:
+            self.entry.runtime_data.resolver_ipv4 = new_resolver
 
     async def async_update(self) -> None:
         """Get the current DNS IP address for hostname."""
-        if self.resolver._closed:  # noqa: SLF001
+        if self._resolver._closed:  # noqa: SLF001
             self.create_dns_resolver()
         response = None
         try:
             async with asyncio.timeout(10):
-                response = await self.resolver.query(self.hostname, self.querytype)
+                response = await self._resolver.query(self.hostname, self.querytype)
         except TimeoutError as err:
             _LOGGER.debug("Timeout while resolving host: %s", err)
-            await self.resolver.close()
+            await self._resolver.close()
         except DNSError as err:
             _LOGGER.warning("Exception while resolving host: %s", err)
-            await self.resolver.close()
+            await self._resolver.close()
 
         if response:
             sorted_ips = sort_ips(