Add coverage for sending answers removes future queued answers (#961)

bdraco · web-flow · commit 2d1b8329ad39 · 2021-08-15T10:15:54.000-05:00
- If we send an answer that is queued to be sent out in the future
  we should remove it from the queue as the question has already
  been answered and we do not want to generate additional traffic.
diff --git a/tests/services/test_info.py b/tests/services/test_info.py
@@ -764,6 +764,7 @@ async def test_we_try_four_times_with_random_delay():
 
     # we are going to patch the zeroconf send to check query transmission
     request_count = 0
+
     def async_send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT):
         """Sends an outgoing packet."""
         nonlocal request_count
diff --git a/tests/test_handlers.py b/tests/test_handlers.py
@@ -1460,3 +1460,42 @@ async def test_response_aggregation_random_delay():
     assert info3.dns_pointer() not in outgoing_queue.queue[1].answers
     assert info4.dns_pointer() not in outgoing_queue.queue[1].answers
     assert info5.dns_pointer() in outgoing_queue.queue[1].answers
+
+
+@pytest.mark.asyncio
+async def test_future_answers_are_removed_on_send():
+    """Verify any future answers scheduled to be sent are removed when we send."""
+    type_ = "_mservice._tcp.local."
+    name = "xxxyyy"
+    registration_name = f"{name}.{type_}"
+
+    desc = {'path': '/~paulsm/'}
+    info = ServiceInfo(
+        type_, registration_name, 80, 0, 0, desc, "ash-1.local.", addresses=[socket.inet_aton("10.0.1.2")]
+    )
+    mocked_zc = unittest.mock.MagicMock()
+    outgoing_queue = MulticastOutgoingQueue(mocked_zc, 0, 0)
+
+    now = current_time_millis()
+    with unittest.mock.patch.object(_handlers, "_MULTICAST_DELAY_RANDOM_INTERVAL", (10, 10)):
+        outgoing_queue.async_add(now, {info.dns_pointer(): set()})
+
+    assert len(outgoing_queue.queue) == 1
+
+    with unittest.mock.patch.object(_handlers, "_MULTICAST_DELAY_RANDOM_INTERVAL", (20, 20)):
+        outgoing_queue.async_add(now, {info.dns_pointer(): set()})
+
+    assert len(outgoing_queue.queue) == 2
+
+    with unittest.mock.patch.object(_handlers, "_MULTICAST_DELAY_RANDOM_INTERVAL", (200, 200)):
+        outgoing_queue.async_add(now, {info.dns_pointer(): set()})
+        outgoing_queue.async_add(now, {info.dns_pointer(): set()})
+
+    assert len(outgoing_queue.queue) == 3
+
+    await asyncio.sleep(0.1)
+    outgoing_queue.async_ready()
+
+    assert len(outgoing_queue.queue) == 1
+    # The answers should all get removed because we just sent them
+    assert len(outgoing_queue.queue[0].answers) == 0
diff --git a/zeroconf/_handlers.py b/zeroconf/_handlers.py
@@ -527,18 +527,24 @@ def async_add(self, now: float, answers: _AnswerWithAdditionalsType) -> None:
                 last_group.answers.update(answers)
                 return
         else:
-            self.zc.loop.call_later(millis_to_seconds(random_delay), self._async_ready)
+            self.zc.loop.call_later(millis_to_seconds(random_delay), self.async_ready)
         self.queue.append(AnswerGroup(send_after, send_before, answers))
 
-    def _async_ready(self) -> None:
+    def _remove_answers_from_queue(self, answers: _AnswerWithAdditionalsType) -> None:
+        """Remove a set of answers from the outgoing queue."""
+        for pending in self.queue:
+            for record in answers:
+                pending.answers.pop(record, None)
+
+    def async_ready(self) -> None:
         """Process anything in the queue that is ready."""
         assert self.zc.loop is not None
         now = current_time_millis()
 
         if len(self.queue) > 1 and self.queue[0].send_before > now:
             # There is more than one answer in the queue,
             # delay until we have to send it (first answer group reaches send_before)
-            self.zc.loop.call_later(millis_to_seconds(self.queue[0].send_before - now), self._async_ready)
+            self.zc.loop.call_later(millis_to_seconds(self.queue[0].send_before - now), self.async_ready)
             return
 
         answers: _AnswerWithAdditionalsType = {}
@@ -549,12 +555,9 @@ def _async_ready(self) -> None:
         if len(self.queue):
             # If there are still groups in the queue that are not ready to send
             # be sure we schedule them to go out later
-            self.zc.loop.call_later(millis_to_seconds(self.queue[0].send_after - now), self._async_ready)
+            self.zc.loop.call_later(millis_to_seconds(self.queue[0].send_after - now), self.async_ready)
 
         if answers:
-            # If we have the same answer scheduled to go out, remove it
-            for pending in self.queue:
-                for record in answers:
-                    pending.answers.pop(record, None)
-
+            # If we have the same answer scheduled to go out, remove them
+            self._remove_answers_from_queue(answers)
             self.zc.async_send(construct_outgoing_multicast_answers(answers))
