Reduce branching in DNSOutgoing.add_answer_at_time (#592)

bdraco · web-flow · commit 35e25fd46f8d · 2021-06-13T11:36:30.000-10:00
diff --git a/tests/test_dns.py b/tests/test_dns.py
@@ -14,7 +14,7 @@
 from typing import Dict, cast  # noqa # used in type hints
 
 import zeroconf as r
-from zeroconf import const
+from zeroconf import const, current_time_millis
 from zeroconf import (
     DNSHinfo,
     DNSText,
@@ -175,6 +175,48 @@ def test_parse_own_packet_response(self):
         assert len(generated.answers) == 1
         assert len(generated.answers) == len(parsed.answers)
 
+    def test_adding_empty_answer(self):
+        generated = r.DNSOutgoing(const._FLAGS_QR_RESPONSE)
+        generated.add_answer_at_time(
+            None,
+            0,
+        )
+        generated.add_answer_at_time(
+            r.DNSService(
+                "æøå.local.",
+                const._TYPE_SRV,
+                const._CLASS_IN | const._CLASS_UNIQUE,
+                const._DNS_HOST_TTL,
+                0,
+                0,
+                80,
+                "foo.local.",
+            ),
+            0,
+        )
+        parsed = r.DNSIncoming(generated.packets()[0])
+        assert len(generated.answers) == 1
+        assert len(generated.answers) == len(parsed.answers)
+
+    def test_adding_expired_answer(self):
+        generated = r.DNSOutgoing(const._FLAGS_QR_RESPONSE)
+        generated.add_answer_at_time(
+            r.DNSService(
+                "æøå.local.",
+                const._TYPE_SRV,
+                const._CLASS_IN | const._CLASS_UNIQUE,
+                const._DNS_HOST_TTL,
+                0,
+                0,
+                80,
+                "foo.local.",
+            ),
+            current_time_millis() + 1000000,
+        )
+        parsed = r.DNSIncoming(generated.packets()[0])
+        assert len(generated.answers) == 0
+        assert len(generated.answers) == len(parsed.answers)
+
     def test_match_question(self):
         generated = r.DNSOutgoing(const._FLAGS_QR_QUERY)
         question = r.DNSQuestion("testname.local.", const._TYPE_SRV, const._CLASS_IN)
diff --git a/zeroconf/_dns.py b/zeroconf/_dns.py
@@ -579,9 +579,8 @@ def add_answer(self, inp: DNSIncoming, record: DNSRecord) -> None:
 
     def add_answer_at_time(self, record: Optional[DNSRecord], now: Union[float, int]) -> None:
         """Adds an answer if it does not expire by a certain time"""
-        if record is not None:
-            if now == 0 or not record.is_expired(now):
-                self.answers.append((record, now))
+        if record is not None and (now == 0 or not record.is_expired(now)):
+            self.answers.append((record, now))
 
     def add_authorative_answer(self, record: DNSPointer) -> None:
         """Adds an authoritative answer"""
