Fix ServiceInfo with multiple A records (#725)

bdraco · web-flow · commit 33385948da91 · 2021-06-18T11:26:15.000-10:00
diff --git a/tests/test_services.py b/tests/test_services.py
@@ -24,6 +24,7 @@
     ServiceInfo,
     ServiceStateChange,
 )
+from zeroconf.aio import AsyncZeroconf
 
 from . import has_working_ipv6, _clear_cache, _inject_response
 
@@ -947,6 +948,28 @@ def test_multiple_addresses():
             assert info.parsed_addresses(r.IPVersion.V6Only) == [address_v6_parsed]
 
 
+# This test uses asyncio because it needs to access the cache directly
+# which is not threadsafe
+@pytest.mark.asyncio
+async def test_multiple_a_addresses():
+    type_ = "_http._tcp.local."
+    registration_name = "multiarec.%s" % type_
+    desc = {'path': '/~paulsm/'}
+    aiozc = AsyncZeroconf(interfaces=['127.0.0.1'])
+    cache = aiozc.zeroconf.cache
+    host = "multahost.local."
+    record1 = r.DNSAddress(host, const._TYPE_A, const._CLASS_IN, 1000, b'a')
+    record2 = r.DNSAddress(host, const._TYPE_A, const._CLASS_IN, 1000, b'b')
+    cache.add(record1)
+    cache.add(record2)
+
+    # New kwarg way
+    info = ServiceInfo(type_, registration_name, 80, 0, 0, desc, host)
+    info.load_from_cache(aiozc.zeroconf)
+    assert set(info.addresses) == set([b'a', b'b'])
+    await aiozc.async_close()
+
+
 def test_backoff():
     got_query = Event()
 
diff --git a/zeroconf/_services/__init__.py b/zeroconf/_services/__init__.py
@@ -777,12 +777,10 @@ def dns_text(self, override_ttl: Optional[int] = None, created: Optional[float]
 
     def _get_address_records_from_cache(self, zc: 'Zeroconf') -> List[DNSRecord]:
         """Get the address records from the cache."""
-        address_records = []
-        cached_a_record = zc.cache.get_by_details(self.server, _TYPE_A, _CLASS_IN)
-        if cached_a_record:
-            address_records.append(cached_a_record)
-        address_records.extend(zc.cache.get_all_by_details(self.server, _TYPE_AAAA, _CLASS_IN))
-        return address_records
+        return [
+            *zc.cache.get_all_by_details(self.server, _TYPE_A, _CLASS_IN),
+            *zc.cache.get_all_by_details(self.server, _TYPE_AAAA, _CLASS_IN),
+        ]
 
     def load_from_cache(self, zc: 'Zeroconf') -> bool:
         """Populate the service info from the cache."""
@@ -844,7 +842,7 @@ def generate_request_query(self, zc: 'Zeroconf', now: float) -> DNSOutgoing:
         out = DNSOutgoing(_FLAGS_QR_QUERY)
         out.add_question_or_one_cache(zc.cache, now, self.name, _TYPE_SRV, _CLASS_IN)
         out.add_question_or_one_cache(zc.cache, now, self.name, _TYPE_TXT, _CLASS_IN)
-        out.add_question_or_one_cache(zc.cache, now, self.server, _TYPE_A, _CLASS_IN)
+        out.add_question_or_all_cache(zc.cache, now, self.server, _TYPE_A, _CLASS_IN)
         out.add_question_or_all_cache(zc.cache, now, self.server, _TYPE_AAAA, _CLASS_IN)
         return out
 
