Run ServiceInfo requests in the event loop (#748)

bdraco · web-flow · commit 0dbcabfade41 · 2021-06-19T09:01:32.000-10:00
diff --git a/examples/async_service_info_request.py b/examples/async_service_info_request.py
@@ -28,7 +28,7 @@ async def async_watch_services(aiozc: AsyncZeroconf) -> None:
             if not name.endswith(HAP_TYPE):
                 continue
             infos.append(AsyncServiceInfo(HAP_TYPE, name))
-        tasks = [info.async_request(aiozc, 3000) for info in infos]
+        tasks = [info.async_request(aiozc.zeroconf, 3000) for info in infos]
         await asyncio.gather(*tasks)
         for info in infos:
             print("Info for %s" % (info.name))
diff --git a/tests/services/test_info.py b/tests/services/test_info.py
@@ -216,7 +216,7 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT):
             send_event.set()
 
         # patch the zeroconf send
-        with unittest.mock.patch.object(zc, "send", send):
+        with unittest.mock.patch.object(zc, "async_send", send):
 
             def mock_incoming_msg(records) -> r.DNSIncoming:
 
@@ -353,7 +353,7 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT):
             send_event.set()
 
         # patch the zeroconf send
-        with unittest.mock.patch.object(zc, "send", send):
+        with unittest.mock.patch.object(zc, "async_send", send):
 
             def mock_incoming_msg(records) -> r.DNSIncoming:
 
diff --git a/tests/test_aio.py b/tests/test_aio.py
@@ -409,7 +409,7 @@ async def test_service_info_async_request() -> None:
     # Generating the race condition is almost impossible
     # without patching since its a TOCTOU race
     with unittest.mock.patch("zeroconf.aio.AsyncServiceInfo._is_complete", False):
-        await aiosinfo.async_request(aiozc, 3000)
+        await aiosinfo.async_request(aiozc.zeroconf, 3000)
     assert aiosinfo is not None
     assert aiosinfo.addresses == [socket.inet_aton("10.0.1.3")]
 
diff --git a/zeroconf/_services/info.py b/zeroconf/_services/info.py
@@ -20,6 +20,7 @@
     USA
 """
 
+import asyncio
 import socket
 from typing import Dict, List, Optional, TYPE_CHECKING, Union, cast
 
@@ -393,6 +394,13 @@ def _is_complete(self) -> bool:
         return not (self.text is None or not self._addresses)
 
     def request(self, zc: 'Zeroconf', timeout: float) -> bool:
+        """Returns true if the service could be discovered on the
+        network, and updates this object with details discovered.
+        """
+        assert zc.loop is not None
+        return asyncio.run_coroutine_threadsafe(self.async_request(zc, timeout), zc.loop).result()
+
+    async def async_request(self, zc: 'Zeroconf', timeout: float) -> bool:
         """Returns true if the service could be discovered on the
         network, and updates this object with details discovered.
         """
@@ -403,22 +411,21 @@ def request(self, zc: 'Zeroconf', timeout: float) -> bool:
         delay = _LISTENER_TIME
         next_ = now
         last = now + timeout
+        await zc.async_wait_for_start()
         try:
-            # Do not set a question on the listener to preload from cache
-            # since we just checked it above in load_from_cache
             zc.add_listener(self, None)
             while not self._is_complete:
                 if last <= now:
                     return False
                 if next_ <= now:
                     out = self.generate_request_query(zc, now)
                     if not out.questions:
-                        return True
-                    zc.send(out)
+                        return self.load_from_cache(zc)
+                    zc.async_send(out)
                     next_ = now + delay
                     delay *= 2
 
-                zc.wait(min(next_, last) - now)
+                await zc.async_wait(min(next_, last) - now)
                 now = current_time_millis()
         finally:
             zc.remove_listener(self)
diff --git a/zeroconf/aio.py b/zeroconf/aio.py
@@ -31,11 +31,10 @@
 from ._services.types import ZeroconfServiceTypes
 from ._utils.aio import wait_condition_or_timeout
 from ._utils.net import IPVersion, InterfaceChoice, InterfacesType
-from ._utils.time import current_time_millis, millis_to_seconds
+from ._utils.time import millis_to_seconds
 from .const import (
     _BROWSER_TIME,
     _CHECK_TIME,
-    _LISTENER_TIME,
     _MDNS_PORT,
     _REGISTER_TIME,
     _SERVICE_TYPE_ENUMERATION_NAME,
@@ -66,38 +65,6 @@ def update_service(self, aiozc: 'AsyncZeroconf', type_: str, name: str) -> None:
 class AsyncServiceInfo(ServiceInfo):
     """An async version of ServiceInfo."""
 
-    async def async_request(self, aiozc: 'AsyncZeroconf', timeout: float) -> bool:
-        """Returns true if the service could be discovered on the
-        network, and updates this object with details discovered.
-        """
-        if self.load_from_cache(aiozc.zeroconf):
-            return True
-
-        now = current_time_millis()
-        delay = _LISTENER_TIME
-        next_ = now
-        last = now + timeout
-        await aiozc.zeroconf.async_wait_for_start()
-        try:
-            aiozc.zeroconf.add_listener(self, None)
-            while not self._is_complete:
-                if last <= now:
-                    return False
-                if next_ <= now:
-                    out = self.generate_request_query(aiozc.zeroconf, now)
-                    if not out.questions:
-                        return self.load_from_cache(aiozc.zeroconf)
-                    aiozc.zeroconf.async_send(out)
-                    next_ = now + delay
-                    delay *= 2
-
-                await aiozc.zeroconf.async_wait(min(next_, last) - now)
-                now = current_time_millis()
-        finally:
-            aiozc.zeroconf.remove_listener(self)
-
-        return True
-
 
 class AsyncServiceBrowser(_ServiceBrowserBase):
     """Used to browse for a service of a specific type.
@@ -333,7 +300,7 @@ async def async_get_service_info(
         name and type, or None if no service matches by the timeout,
         which defaults to 3 seconds."""
         info = AsyncServiceInfo(type_, name)
-        if await info.async_request(self, timeout):
+        if await info.async_request(self.zeroconf, timeout):
             return info
         return None
 
