From 19b301aa35350d7db76e7f38f895b9d479791e8d Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 27 Jun 2021 07:48:12 -1000
Subject: [PATCH] Make ServiceInfo first question QU

- We want an immediate response when making a request with ServiceInfo
  by asking a QU question, most responders will not delay the response
  and respond right away to our question. This also improves compatibility
  with split networks as we may not have been able to see the response
  otherwise.
---
 tests/services/test_info.py |  6 ++---
 tests/test_aio.py           | 47 +++++++++++++++++++++++++++++++++++++
 zeroconf/_services/info.py  |  6 ++++-
 3 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/tests/services/test_info.py b/tests/services/test_info.py
index adca1a53..37f98aa1 100644
--- a/tests/services/test_info.py
+++ b/tests/services/test_info.py
@@ -682,8 +682,8 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT):
         zeroconf.close()
 
 
-def test_asking_qm_questions_are_default():
-    """Verify default is QM questions."""
+def test_asking_qm_questions():
+    """Verify explictly asking QM questions."""
     type_ = "_quservice._tcp.local."
     zeroconf = r.Zeroconf(interfaces=['127.0.0.1'])
 
@@ -701,6 +701,6 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT):
 
     # patch the zeroconf send
     with patch.object(zeroconf, "async_send", send):
-        zeroconf.get_service_info(f"name.{type_}", type_, 500)
+        zeroconf.get_service_info(f"name.{type_}", type_, 500, question_type=r.DNSQuestionType.QM)
         assert first_outgoing.questions[0].unicast == False
         zeroconf.close()
diff --git a/tests/test_aio.py b/tests/test_aio.py
index 41e0e83a..f22bf966 100644
--- a/tests/test_aio.py
+++ b/tests/test_aio.py
@@ -776,3 +776,50 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT, v6_flow_scope=()):
             assert service_removed.is_set()
             await browser.async_cancel()
             await aiozc.async_close()
+
+
+@pytest.mark.asyncio
+async def test_info_asking_default_is_asking_qm_questions_after_the_first_qu():
+    """Verify the service info first question is QU and subsequent ones are QM questions."""
+    type_ = "_quservice._tcp.local."
+    aiozc = AsyncZeroconf(interfaces=['127.0.0.1'])
+    zeroconf_info = aiozc.zeroconf
+
+    name = "xxxyyy"
+    registration_name = "%s.%s" % (name, type_)
+
+    desc = {'path': '/~paulsm/'}
+    info = ServiceInfo(
+        type_, registration_name, 80, 0, 0, desc, "ash-2.local.", addresses=[socket.inet_aton("10.0.1.2")]
+    )
+
+    zeroconf_info.registry.add(info)
+
+    # we are going to patch the zeroconf send to check query transmission
+    old_send = zeroconf_info.async_send
+
+    first_outgoing = None
+    second_outgoing = None
+
+    def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT):
+        """Sends an outgoing packet."""
+        nonlocal first_outgoing
+        nonlocal second_outgoing
+        if out.questions:
+            if first_outgoing is not None and second_outgoing is None:
+                second_outgoing = out
+            if first_outgoing is None:
+                first_outgoing = out
+        old_send(out, addr=addr, port=port)
+
+    # patch the zeroconf send
+    with patch.object(zeroconf_info, "async_send", send):
+        aiosinfo = AsyncServiceInfo(type_, registration_name)
+        # Patch _is_complete so we send multiple times
+        with patch("zeroconf.aio.AsyncServiceInfo._is_complete", False):
+            await aiosinfo.async_request(aiozc.zeroconf, 1200)
+        try:
+            assert first_outgoing.questions[0].unicast == True
+            assert second_outgoing.questions[0].unicast == False
+        finally:
+            await aiozc.async_close()
diff --git a/zeroconf/_services/info.py b/zeroconf/_services/info.py
index 4365d6ef..9d1c37f3 100644
--- a/zeroconf/_services/info.py
+++ b/zeroconf/_services/info.py
@@ -438,6 +438,7 @@ async def async_request(
         if self.load_from_cache(zc):
             return True
 
+        first_request = True
         now = current_time_millis()
         delay = _LISTENER_TIME
         next_ = now
@@ -449,7 +450,10 @@ async def async_request(
                 if last <= now:
                     return False
                 if next_ <= now:
-                    out = self.generate_request_query(zc, now, question_type)
+                    out = self.generate_request_query(
+                        zc, now, question_type or DNSQuestionType.QU if first_request else DNSQuestionType.QM
+                    )
+                    first_request = False
                     if not out.questions:
                         return self.load_from_cache(zc)
                     zc.async_send(out)