diff --git a/homeassistant/components/dnsip/__init__.py b/homeassistant/components/dnsip/__init__.py index ec5a9f033d248..69c385643f15a 100644 --- a/homeassistant/components/dnsip/__init__.py +++ b/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: diff --git a/homeassistant/components/dnsip/sensor.py b/homeassistant/components/dnsip/sensor.py index dd5a4f38aab37..ca965f330195e 100644 --- a/homeassistant/components/dnsip/sensor.py +++ b/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,7 +46,7 @@ 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.""" @@ -54,16 +54,29 @@ async def async_setup_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,10 +88,9 @@ 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, @@ -86,6 +98,8 @@ def __init__( 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( diff --git a/tests/components/dnsip/test_init.py b/tests/components/dnsip/test_init.py index 1181c391ca2f5..db32a22895818 100644 --- a/tests/components/dnsip/test_init.py +++ b/tests/components/dnsip/test_init.py @@ -1,7 +1,12 @@ """Test for DNS IP integration Init.""" +import asyncio from unittest.mock import patch +from aiodns.error import DNSError +from pycares import AresError +import pytest + from homeassistant.components.dnsip.const import ( CONF_HOSTNAME, CONF_IPV4, @@ -44,7 +49,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,8 +87,8 @@ async def test_port_migration( entry.add_to_hass(hass) with patch( - "homeassistant.components.dnsip.sensor.aiodns.DNSResolver", - return_value=RetrieveDNS(), + "homeassistant.components.dnsip.aiodns.DNSResolver", + side_effect=[RetrieveDNS(), RetrieveDNS()], ): await hass.config_entries.async_setup(entry.entry_id) await hass.async_block_till_done() @@ -123,8 +128,8 @@ async def test_remove_unique_id_migration( entry.add_to_hass(hass) with patch( - "homeassistant.components.dnsip.sensor.aiodns.DNSResolver", - return_value=RetrieveDNS(), + "homeassistant.components.dnsip.aiodns.DNSResolver", + side_effect=[RetrieveDNS(), RetrieveDNS()], ): await hass.config_entries.async_setup(entry.entry_id) await hass.async_block_till_done() @@ -160,7 +165,7 @@ 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) @@ -168,3 +173,119 @@ async def test_migrate_error_from_future(hass: HomeAssistant) -> None: entry = hass.config_entries.async_get_entry(entry.entry_id) assert entry.state is ConfigEntryState.MIGRATION_ERROR + + +@pytest.mark.parametrize( + "error", + [ + TimeoutError(), + DNSError(), + AresError(), + asyncio.CancelledError(), + ], +) +async def test_setup_dns_error(hass: HomeAssistant, error: Exception) -> 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=error), + ): + 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 creates only the IPv6 entity.""" + + 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 + assert hass.states.get("sensor.home_assistant_io_ipv6") is not None + assert hass.states.get("sensor.home_assistant_io") is None + + +async def test_setup_dns_timeout(hass: HomeAssistant) -> None: + """Test setup raises ConfigEntryNotReady when DNS lookup times out.""" + + 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(), + ), + patch( + "homeassistant.components.dnsip.asyncio.timeout", + side_effect=TimeoutError(), + ), + ): + await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() + + assert entry.state is ConfigEntryState.SETUP_RETRY diff --git a/tests/components/dnsip/test_sensor.py b/tests/components/dnsip/test_sensor.py index b7dae63476511..ef8f7809dd641 100644 --- a/tests/components/dnsip/test_sensor.py +++ b/tests/components/dnsip/test_sensor.py @@ -48,8 +48,8 @@ async def test_sensor(hass: HomeAssistant) -> None: entry.add_to_hass(hass) with patch( - "homeassistant.components.dnsip.sensor.aiodns.DNSResolver", - return_value=RetrieveDNS(), + "homeassistant.components.dnsip.aiodns.DNSResolver", + side_effect=[RetrieveDNS(), RetrieveDNS()], ): await hass.config_entries.async_setup(entry.entry_id) await hass.async_block_till_done() @@ -91,8 +91,8 @@ async def test_legacy_sensor(hass: HomeAssistant) -> None: entry.add_to_hass(hass) with patch( - "homeassistant.components.dnsip.sensor.aiodns.DNSResolver", - return_value=RetrieveDNS(), + "homeassistant.components.dnsip.aiodns.DNSResolver", + side_effect=[RetrieveDNS(), RetrieveDNS()], ): await hass.config_entries.async_setup(entry.entry_id) await hass.async_block_till_done() @@ -135,10 +135,10 @@ async def test_sensor_no_response( ) entry.add_to_hass(hass) - dns_mock = RetrieveDNS() + dns_mock_ipv4 = RetrieveDNS() with patch( - "homeassistant.components.dnsip.sensor.aiodns.DNSResolver", - return_value=dns_mock, + "homeassistant.components.dnsip.aiodns.DNSResolver", + return_value=dns_mock_ipv4, ): await hass.config_entries.async_setup(entry.entry_id) await hass.async_block_till_done() @@ -147,10 +147,10 @@ async def test_sensor_no_response( assert state.state == "1.1.1.1" - dns_mock.error = DNSError() + dns_mock_ipv4.error = DNSError() with patch( "homeassistant.components.dnsip.sensor.aiodns.DNSResolver", - return_value=dns_mock, + return_value=dns_mock_ipv4, ): freezer.tick(timedelta(seconds=SCAN_INTERVAL.seconds)) async_fire_time_changed(hass) @@ -158,7 +158,6 @@ async def test_sensor_no_response( 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"] @@ -195,10 +194,10 @@ async def test_sensor_timeout( ) entry.add_to_hass(hass) - dns_mock = RetrieveDNS() + dns_mock_ipv4 = RetrieveDNS() with patch( - "homeassistant.components.dnsip.sensor.aiodns.DNSResolver", - return_value=dns_mock, + "homeassistant.components.dnsip.aiodns.DNSResolver", + return_value=dns_mock_ipv4, ): await hass.config_entries.async_setup(entry.entry_id) await hass.async_block_till_done() @@ -210,7 +209,7 @@ async def test_sensor_timeout( with ( patch( "homeassistant.components.dnsip.sensor.aiodns.DNSResolver", - return_value=dns_mock, + return_value=dns_mock_ipv4, ), patch( "homeassistant.components.dnsip.sensor.asyncio.timeout", @@ -221,7 +220,6 @@ async def test_sensor_timeout( 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"]