Use runtime_data and validate connection at setup for dnsip

homeassistant/components/dnsip/__init__.py

@@ -1,23 +1,83 @@
 """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."""
 
+    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]
+    try:
+        async with asyncio.timeout(10):
+            if entry.data[CONF_IPV4]:
+                await resolver_ipv4.query(hostname, "A")
+            elif entry.data[CONF_IPV6]:
+                await resolver_ipv6.query(hostname, "AAAA")
+    except (TimeoutError, DNSError) as err:
+        await resolver_ipv4.close()
+        await resolver_ipv6.close()
+        raise ConfigEntryNotReady(f"DNS lookup failed for {hostname}: {err}") from err
+
+    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:
+        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:

homeassistant/components/dnsip/sensor.py

@@ -10,17 +10,16 @@
 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.const import CONF_NAME
 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,
     CONF_IPV6,
-    CONF_PORT_IPV6,
     CONF_RESOLVER,
     CONF_RESOLVER_IPV6,
     DOMAIN,
@@ -46,24 +45,35 @@ 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.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.runtime_data.resolver_ipv6,
+            )
+        )
 
     async_add_entities(entities, update_before_add=True)
 
@@ -75,22 +85,19 @@ 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,28 +111,17 @@ def __init__(
             model=aiodns.__version__,
             name=name,
         )
-        self.create_dns_resolver()
-
-    def create_dns_resolver(self) -> None:
-        """Create the DNS resolver."""
-        self.resolver = aiodns.DNSResolver(
-            nameservers=[self.nameserver], tcp_port=self.port, udp_port=self.port
-        )
 
     async def async_update(self) -> None:
         """Get the current DNS IP address for hostname."""
-        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)
         except TimeoutError as err:
             _LOGGER.debug("Timeout while resolving host: %s", err)
-            await self.resolver.close()
         except DNSError as err:
             _LOGGER.warning("Exception while resolving host: %s", err)
-            await self.resolver.close()
 
         if response:
             sorted_ips = sort_ips(

tests/components/dnsip/test_init.py

@@ -2,6 +2,8 @@
 
 from unittest.mock import patch
 
+from aiodns.error import DNSError
+
 from homeassistant.components.dnsip.const import (
     CONF_HOSTNAME,
     CONF_IPV4,
@@ -44,7 +46,7 @@ async def test_load_unload_entry(hass: HomeAssistant) -> None:
     entry.add_to_hass(hass)
 
     with patch(
-        "homeassistant.components.dnsip.config_flow.aiodns.DNSResolver",
+        "homeassistant.components.dnsip.aiodns.DNSResolver",
         return_value=RetrieveDNS(),
     ):
         await hass.config_entries.async_setup(entry.entry_id)
@@ -82,7 +84,7 @@ async def test_port_migration(
     entry.add_to_hass(hass)
 
     with patch(
-        "homeassistant.components.dnsip.sensor.aiodns.DNSResolver",
+        "homeassistant.components.dnsip.aiodns.DNSResolver",
         return_value=RetrieveDNS(),
     ):
         await hass.config_entries.async_setup(entry.entry_id)
@@ -123,7 +125,7 @@ async def test_remove_unique_id_migration(
     entry.add_to_hass(hass)
 
     with patch(
-        "homeassistant.components.dnsip.sensor.aiodns.DNSResolver",
+        "homeassistant.components.dnsip.aiodns.DNSResolver",
         return_value=RetrieveDNS(),
     ):
         await hass.config_entries.async_setup(entry.entry_id)
@@ -160,11 +162,77 @@ async def test_migrate_error_from_future(hass: HomeAssistant) -> None:
     entry.add_to_hass(hass)
 
     with patch(
-        "homeassistant.components.dnsip.sensor.aiodns.DNSResolver",
+        "homeassistant.components.dnsip.aiodns.DNSResolver",
         return_value=RetrieveDNS(),
     ):
         await hass.config_entries.async_setup(entry.entry_id)
         await hass.async_block_till_done()
 
     entry = hass.config_entries.async_get_entry(entry.entry_id)
     assert entry.state is ConfigEntryState.MIGRATION_ERROR
+
+
+async def test_setup_dns_error(hass: HomeAssistant) -> None:
+    """Test setup raises ConfigEntryNotReady when DNS lookup fails."""
+
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        source=SOURCE_USER,
+        data={
+            CONF_HOSTNAME: "home-assistant.io",
+            CONF_NAME: "home-assistant.io",
+            CONF_IPV4: True,
+            CONF_IPV6: False,
+        },
+        options={
+            CONF_RESOLVER: "208.67.222.222",
+            CONF_RESOLVER_IPV6: "2620:119:53::53",
+            CONF_PORT: 53,
+            CONF_PORT_IPV6: 53,
+        },
+        entry_id="1",
+        unique_id="home-assistant.io",
+    )
+    entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.dnsip.aiodns.DNSResolver",
+        return_value=RetrieveDNS(error=DNSError()),
+    ):
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_setup_ipv6_only(hass: HomeAssistant) -> None:
+    """Test setup with only IPv6 enabled exercises the IPv6 lookup branch."""
+
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        source=SOURCE_USER,
+        data={
+            CONF_HOSTNAME: "home-assistant.io",
+            CONF_NAME: "home-assistant.io",
+            CONF_IPV4: False,
+            CONF_IPV6: True,
+        },
+        options={
+            CONF_RESOLVER: "208.67.222.222",
+            CONF_RESOLVER_IPV6: "2620:119:53::53",
+            CONF_PORT: 53,
+            CONF_PORT_IPV6: 53,
+        },
+        entry_id="1",
+        unique_id="home-assistant.io",
+    )
+    entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.dnsip.aiodns.DNSResolver",
+        return_value=RetrieveDNS(),
+    ):
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.LOADED