Add timeout to dnsip (to handle stale connections) (home-assistant#153086)

Author: gjohansson-ST

homeassistant/components/dnsip/sensor.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from datetime import timedelta
 from ipaddress import IPv4Address, IPv6Address
 import logging
@@ -88,8 +89,8 @@ def __init__(
         self._attr_name = "IPv6" if ipv6 else None
         self._attr_unique_id = f"{hostname}_{ipv6}"
         self.hostname = hostname
-        self.resolver = aiodns.DNSResolver(tcp_port=port, udp_port=port)
-        self.resolver.nameservers = [resolver]
+        self.port = port
+        self._resolver = resolver
         self.querytype: Literal["A", "AAAA"] = "AAAA" if ipv6 else "A"
         self._retries = DEFAULT_RETRIES
         self._attr_extra_state_attributes = {
@@ -103,14 +104,26 @@ def __init__(
             model=aiodns.__version__,
             name=name,
         )
+        self.resolver: aiodns.DNSResolver
+        self.create_dns_resolver()
+
+    def create_dns_resolver(self) -> None:
+        """Create the DNS resolver."""
+        self.resolver = aiodns.DNSResolver(tcp_port=self.port, udp_port=self.port)
+        self.resolver.nameservers = [self._resolver]
 
     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:
-            response = await self.resolver.query(self.hostname, self.querytype)
+            async with asyncio.timeout(10):
+                response = await self.resolver.query(self.hostname, self.querytype)
+        except TimeoutError:
+            await self.resolver.close()
         except DNSError as err:
             _LOGGER.warning("Exception while resolving host: %s", err)
-            response = None
 
         if response:
             sorted_ips = sort_ips(

tests/components/dnsip/__init__.py

@@ -23,6 +23,7 @@ def __init__(
             self.nameservers = nameservers
         self._nameservers = ["1.2.3.4"]
         self.error = error
+        self._closed = False
 
     async def query(self, hostname, qtype) -> list[QueryResult]:
         """Return information."""
@@ -47,3 +48,7 @@ def nameservers(self) -> list[str]:
     @nameservers.setter
     def nameservers(self, value: list[str]) -> None:
         self._nameservers = value
+
+    async def close(self) -> None:
+        """Close the resolver."""
+        self._closed = True

tests/components/dnsip/test_sensor.py

@@ -171,3 +171,70 @@ async def test_sensor_no_response(
 
     state = hass.states.get("sensor.home_assistant_io")
     assert state.state == STATE_UNAVAILABLE
+
+
+async def test_sensor_timeout(
+    hass: HomeAssistant, freezer: FrozenDateTimeFactory
+) -> None:
+    """Test the DNS IP sensor with timeout."""
+    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)
+
+    dns_mock = RetrieveDNS()
+    with patch(
+        "homeassistant.components.dnsip.sensor.aiodns.DNSResolver",
+        return_value=dns_mock,
+    ):
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.home_assistant_io")
+
+    assert state.state == "1.1.1.1"
+
+    with (
+        patch(
+            "homeassistant.components.dnsip.sensor.aiodns.DNSResolver",
+            return_value=dns_mock,
+        ),
+        patch(
+            "homeassistant.components.dnsip.sensor.asyncio.timeout",
+            side_effect=TimeoutError(),
+        ),
+    ):
+        freezer.tick(timedelta(seconds=SCAN_INTERVAL.seconds))
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+
+        # Allows 2 retries before going unavailable
+        state = hass.states.get("sensor.home_assistant_io")
+        assert state.state == "1.1.1.1"
+        assert state.attributes["ip_addresses"] == ["1.1.1.1", "1.2.3.4"]
+
+        freezer.tick(timedelta(seconds=SCAN_INTERVAL.seconds))
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+
+        freezer.tick(timedelta(seconds=SCAN_INTERVAL.seconds))
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.home_assistant_io")
+    assert state.state == STATE_UNAVAILABLE