From 1c414966fe3c5d9f79254796a4661e468b1d8ea4 Mon Sep 17 00:00:00 2001 From: pemontto <939704+pemontto@users.noreply.github.com> Date: Tue, 7 May 2024 10:49:13 +0100 Subject: [PATCH] Add support for round-robin DNS (#115218) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add support for RR DNS * ๐Ÿงช Update tests for DNS IP round-robin * ๐Ÿค– Configure DNS IP round-robin automatically * ๐Ÿ› Sort IPv6 addresses correctly * Limit returned IPs and cleanup test class * ๐Ÿ”Ÿ Change max DNS results to 10 * Rename IPs to ip_addresses --- homeassistant/components/dnsip/config_flow.py | 5 ++++- homeassistant/components/dnsip/sensor.py | 19 ++++++++++++++++++- tests/components/dnsip/__init__.py | 19 +++++++++++++++---- tests/components/dnsip/test_sensor.py | 16 ++++++++++++---- 4 files changed, 49 insertions(+), 10 deletions(-) diff --git a/homeassistant/components/dnsip/config_flow.py b/homeassistant/components/dnsip/config_flow.py index f07971d5db5..21a29465050 100644 --- a/homeassistant/components/dnsip/config_flow.py +++ b/homeassistant/components/dnsip/config_flow.py @@ -176,7 +176,10 @@ class DnsIPOptionsFlowHandler(OptionsFlowWithConfigEntry): else: return self.async_create_entry( title=self.config_entry.title, - data={CONF_RESOLVER: resolver, CONF_RESOLVER_IPV6: resolver_ipv6}, + data={ + CONF_RESOLVER: resolver, + CONF_RESOLVER_IPV6: resolver_ipv6, + }, ) schema = self.add_suggested_values_to_schema( diff --git a/homeassistant/components/dnsip/sensor.py b/homeassistant/components/dnsip/sensor.py index 529de6f2b1b..d3527bda3f2 100644 --- a/homeassistant/components/dnsip/sensor.py +++ b/homeassistant/components/dnsip/sensor.py @@ -3,6 +3,7 @@ from __future__ import annotations from datetime import timedelta +from ipaddress import IPv4Address, IPv6Address import logging import aiodns @@ -25,12 +26,23 @@ from .const import ( ) DEFAULT_RETRIES = 2 +MAX_RESULTS = 10 _LOGGER = logging.getLogger(__name__) SCAN_INTERVAL = timedelta(seconds=120) +def sort_ips(ips: list, querytype: str) -> list: + """Join IPs into a single string.""" + + if querytype == "AAAA": + ips = [IPv6Address(ip) for ip in ips] + else: + ips = [IPv4Address(ip) for ip in ips] + return [str(ip) for ip in sorted(ips)][:MAX_RESULTS] + + async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: @@ -41,6 +53,7 @@ async def async_setup_entry( resolver_ipv4 = entry.options[CONF_RESOLVER] resolver_ipv6 = entry.options[CONF_RESOLVER_IPV6] + entities = [] if entry.data[CONF_IPV4]: entities.append(WanIpSensor(name, hostname, resolver_ipv4, False)) @@ -92,7 +105,11 @@ class WanIpSensor(SensorEntity): response = None if response: - self._attr_native_value = response[0].host + sorted_ips = sort_ips( + [res.host for res in response], querytype=self.querytype + ) + self._attr_native_value = sorted_ips[0] + self._attr_extra_state_attributes["ip_addresses"] = sorted_ips self._attr_available = True self._retries = DEFAULT_RETRIES elif self._retries > 0: diff --git a/tests/components/dnsip/__init__.py b/tests/components/dnsip/__init__.py index d98de181892..a0e6b7c81b8 100644 --- a/tests/components/dnsip/__init__.py +++ b/tests/components/dnsip/__init__.py @@ -6,8 +6,10 @@ from __future__ import annotations class QueryResult: """Return Query results.""" - host = "1.2.3.4" - ttl = 60 + def __init__(self, ip="1.2.3.4", ttl=60) -> None: + """Initialize QueryResult class.""" + self.host = ip + self.ttl = ttl class RetrieveDNS: @@ -22,11 +24,20 @@ class RetrieveDNS: self._nameservers = ["1.2.3.4"] self.error = error - async def query(self, hostname, qtype) -> dict[str, str]: + async def query(self, hostname, qtype) -> list[QueryResult]: """Return information.""" if self.error: raise self.error - return [QueryResult] + if qtype == "AAAA": + results = [ + QueryResult("2001:db8:77::face:b00c"), + QueryResult("2001:db8:77::dead:beef"), + QueryResult("2001:db8::77:dead:beef"), + QueryResult("2001:db8:66::dead:beef"), + ] + else: + results = [QueryResult("1.2.3.4"), QueryResult("1.1.1.1")] + return results @property def nameservers(self) -> list[str]: diff --git a/tests/components/dnsip/test_sensor.py b/tests/components/dnsip/test_sensor.py index e1353d83268..0a81804a689 100644 --- a/tests/components/dnsip/test_sensor.py +++ b/tests/components/dnsip/test_sensor.py @@ -56,8 +56,15 @@ async def test_sensor(hass: HomeAssistant) -> None: state1 = hass.states.get("sensor.home_assistant_io") state2 = hass.states.get("sensor.home_assistant_io_ipv6") - assert state1.state == "1.2.3.4" - assert state2.state == "1.2.3.4" + assert state1.state == "1.1.1.1" + assert state1.attributes["ip_addresses"] == ["1.1.1.1", "1.2.3.4"] + assert state2.state == "2001:db8::77:dead:beef" + assert state2.attributes["ip_addresses"] == [ + "2001:db8::77:dead:beef", + "2001:db8:66::dead:beef", + "2001:db8:77::dead:beef", + "2001:db8:77::face:b00c", + ] async def test_sensor_no_response( @@ -92,7 +99,7 @@ async def test_sensor_no_response( state = hass.states.get("sensor.home_assistant_io") - assert state.state == "1.2.3.4" + assert state.state == "1.1.1.1" dns_mock.error = DNSError() with patch( @@ -107,7 +114,8 @@ async def test_sensor_no_response( # Allows 2 retries before going unavailable state = hass.states.get("sensor.home_assistant_io") - assert state.state == "1.2.3.4" + 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)