Add coverage for sending answers removes future queued answers

bdraco · bdraco · commit ed52266b553e · 2021-08-15T10:04:33.000-05:00
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))
