Add an AsyncZeroconfServiceTypes to mirror ZeroconfServiceTypes to zeroconf.aio (#658)

bdraco · web-flow · commit aaf8a368063f · 2021-06-15T14:04:06.000-10:00
diff --git a/tests/test_aio.py b/tests/test_aio.py
@@ -5,13 +5,14 @@
 """Unit tests for aio.py."""
 
 import asyncio
+import logging
 import socket
 import threading
 import unittest.mock
 
 import pytest
 
-from zeroconf.aio import AsyncServiceInfo, AsyncServiceListener, AsyncZeroconf
+from zeroconf.aio import AsyncServiceInfo, AsyncServiceListener, AsyncZeroconf, AsyncZeroconfServiceTypes
 from zeroconf import Zeroconf
 from zeroconf.const import _LISTENER_TIME
 from zeroconf._exceptions import BadTypeInNameException, NonUniqueNameException, ServiceNameAlreadyRegistered
@@ -20,8 +21,19 @@
 
 from . import _clear_cache
 
+log = logging.getLogger('zeroconf')
+original_logging_level = logging.NOTSET
+
+
+def setup_module():
+    global original_logging_level
+    original_logging_level = log.level
+    log.setLevel(logging.DEBUG)
 
-from . import _clear_cache
+
+def teardown_module():
+    if original_logging_level != logging.NOTSET:
+        log.setLevel(original_logging_level)
 
 
 @pytest.fixture(autouse=True)
@@ -558,3 +570,37 @@ async def test_async_unregister_all_services() -> None:
     await aiozc.async_unregister_all_services()
 
     await aiozc.async_close()
+
+
+@pytest.mark.asyncio
+async def test_async_zeroconf_service_types():
+    type_ = "_test-srvc-type._tcp.local."
+    name = "xxxyyy"
+    registration_name = "%s.%s" % (name, type_)
+
+    zeroconf_registrar = AsyncZeroconf(interfaces=['127.0.0.1'])
+    desc = {'path': '/~paulsm/'}
+    info = ServiceInfo(
+        type_,
+        registration_name,
+        80,
+        0,
+        0,
+        desc,
+        "ash-2.local.",
+        addresses=[socket.inet_aton("10.0.1.2")],
+    )
+    task = await zeroconf_registrar.async_register_service(info)
+    await task
+    # Ensure we do not clear the cache until after the last broadcast is processed
+    await asyncio.sleep(0.2)
+    _clear_cache(zeroconf_registrar.zeroconf)
+    try:
+        service_types = await AsyncZeroconfServiceTypes.async_find(interfaces=['127.0.0.1'], timeout=0.5)
+        assert type_ in service_types
+        _clear_cache(zeroconf_registrar.zeroconf)
+        service_types = await AsyncZeroconfServiceTypes.async_find(aiozc=zeroconf_registrar, timeout=0.5)
+        assert type_ in service_types
+
+    finally:
+        await zeroconf_registrar.async_close()
diff --git a/zeroconf/aio.py b/zeroconf/aio.py
@@ -22,22 +22,32 @@
 import asyncio
 import contextlib
 from types import TracebackType  # noqa # used in type hints
-from typing import Awaitable, Callable, Dict, List, Optional, Type, Union
+from typing import Awaitable, Callable, Dict, List, Optional, Tuple, Type, Union
 
 from ._core import NotifyListener, Zeroconf
 from ._exceptions import NonUniqueNameException
 from ._services import ServiceInfo, _ServiceBrowserBase, instance_name_from_service_info
+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 .const import _BROWSER_TIME, _CHECK_TIME, _LISTENER_TIME, _MDNS_PORT, _REGISTER_TIME, _UNREGISTER_TIME
+from .const import (
+    _BROWSER_TIME,
+    _CHECK_TIME,
+    _LISTENER_TIME,
+    _MDNS_PORT,
+    _REGISTER_TIME,
+    _SERVICE_TYPE_ENUMERATION_NAME,
+    _UNREGISTER_TIME,
+)
 
 
 __all__ = [
     "AsyncZeroconf",
     "AsyncServiceInfo",
     "AsyncServiceBrowser",
     "AsyncServiceListener",
+    "AsyncZeroconfServiceTypes",
 ]
 
 
@@ -137,6 +147,7 @@ async def async_cancel(self) -> None:
     async def async_run(self) -> None:
         """Run the browser task."""
         self.run()
+        await self.aiozc.zeroconf.async_wait_for_start()
         while True:
             timeout = self._seconds_to_wait()
             if timeout:
@@ -164,6 +175,45 @@ async def async_run(self) -> None:
             )
 
 
+class AsyncZeroconfServiceTypes(ZeroconfServiceTypes):
+    """An async version of ZeroconfServiceTypes."""
+
+    @classmethod
+    async def async_find(
+        cls,
+        aiozc: Optional['AsyncZeroconf'] = None,
+        timeout: Union[int, float] = 5,
+        interfaces: InterfacesType = InterfaceChoice.All,
+        ip_version: Optional[IPVersion] = None,
+    ) -> Tuple[str, ...]:
+        """
+        Return all of the advertised services on any local networks.
+
+        :param aiozc: AsyncZeroconf() instance.  Pass in if already have an
+                instance running or if non-default interfaces are needed
+        :param timeout: seconds to wait for any responses
+        :param interfaces: interfaces to listen on.
+        :param ip_version: IP protocol version to use.
+        :return: tuple of service type strings
+        """
+        local_zc = aiozc or AsyncZeroconf(interfaces=interfaces, ip_version=ip_version)
+        listener = cls()
+        async_browser = AsyncServiceBrowser(
+            local_zc, _SERVICE_TYPE_ENUMERATION_NAME, listener=listener  # type: ignore
+        )
+
+        # wait for responses
+        await asyncio.sleep(timeout)
+
+        await async_browser.async_cancel()
+
+        # close down anything we opened
+        if aiozc is None:
+            await local_zc.async_close()
+
+        return tuple(sorted(listener.found_services))
+
+
 class AsyncZeroconf:
     """Implementation of Zeroconf Multicast DNS Service Discovery
 
