Use runtime_data and validate connection at setup for dnsip

homeassistant/components/dnsip/__init__.py

@@ -1,26 +1,109 @@
 """The DNS IP integration."""
 
+import asyncio
+from dataclasses import dataclass
+
+import aiodns
+from aiodns.error import DNSError
+
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PORT
 from homeassistant.core import _LOGGER, 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,
+)
+
 
-from .const import CONF_PORT_IPV6, DEFAULT_PORT, PLATFORMS
+@dataclass
+class DnsIPRuntimeData:
+    """Runtime data for DNS IP integration."""
 
+    resolver_ipv4: aiodns.DNSResolver
+    resolver_ipv6: aiodns.DNSResolver
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+type DnsIPConfigEntry = ConfigEntry[DnsIPRuntimeData]
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: DnsIPConfigEntry) -> bool:
     """Set up DNS IP from a config entry."""
 
-    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+    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]
+
+    resolver_ipv4 = aiodns.DNSResolver(
+        nameservers=[nameserver_ipv4], tcp_port=port_ipv4, udp_port=port_ipv4
+    )
+    resolver_ipv6 = aiodns.DNSResolver(
+        nameservers=[nameserver_ipv6], tcp_port=port_ipv6, udp_port=port_ipv6
+    )
+
+    hostname = entry.data[CONF_HOSTNAME]
+    queries: list = []
+    if entry.data[CONF_IPV4]:
+        queries.append(resolver_ipv4.query(hostname, "A"))
+    if entry.data[CONF_IPV6]:
+        queries.append(resolver_ipv6.query(hostname, "AAAA"))
+
+    try:
+        async with asyncio.timeout(10):
+            results = await asyncio.gather(*queries, return_exceptions=True)
+    except TimeoutError as err:
+        await resolver_ipv4.close()
+        await resolver_ipv6.close()
+        raise ConfigEntryNotReady(
+            f"DNS lookup timed out for {hostname}: {err}"
+        ) from err
+
+    errors = [
+        result for result in results if isinstance(result, (TimeoutError, DNSError))
+    ]
+    if errors and len(errors) == len(results):
+        await resolver_ipv4.close()
+        await resolver_ipv6.close()
+        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,
+    )
+
+    try:
+        await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+    except Exception:
+        await resolver_ipv4.close()
+        await resolver_ipv6.close()
+        raise
+
     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:
+        await entry.runtime_data.resolver_ipv4.close()
+        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

@@ -10,12 +10,12 @@
 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(
+                name,
+                hostname,
+                entry.options[CONF_RESOLVER],
+                False,
+                entry.options[CONF_PORT],
+                entry.runtime_data.resolver_ipv4,
+            )
+        )
     if entry.data[CONF_IPV6]:
-        entities.append(WanIpSensor(name, hostname, nameserver_ipv6, True, port_ipv6))
+        entities.append(
+            WanIpSensor(
+                name,
+                hostname,
+                entry.options[CONF_RESOLVER_IPV6],
+                True,
+                entry.options[CONF_PORT_IPV6],
+                entry.runtime_data.resolver_ipv6,
+            )
+        )
 
     async_add_entities(entities, update_before_add=True)
 
@@ -75,22 +88,22 @@ class WanIpSensor(SensorEntity):
     _attr_translation_key = "dnsip"
     _unrecorded_attributes = frozenset({"resolver", "querytype", "ip_addresses"})
 
-    resolver: aiodns.DNSResolver
-
     def __init__(
         self,
         name: str,
         hostname: str,
         nameserver: str,
         ipv6: bool,
         port: int,
+        resolver: aiodns.DNSResolver,
     ) -> None:
         """Initialize the DNS IP sensor."""
         self._attr_name = "IPv6" if ipv6 else None
         self._attr_unique_id = f"{hostname}_{ipv6}"
         self.hostname = hostname
         self.port = port
         self.nameserver = nameserver
+        self.resolver = resolver
         self.querytype: Literal["A", "AAAA"] = "AAAA" if ipv6 else "A"
         self._retries = DEFAULT_RETRIES
         self._attr_extra_state_attributes = {
@@ -104,7 +117,6 @@ def __init__(
             model=aiodns.__version__,
             name=name,
         )
-        self.create_dns_resolver()
 
     def create_dns_resolver(self) -> None:
         """Create the DNS resolver."""