Add async_unregister_all_services to AsyncZeroconf (#649)

bdraco · web-flow · commit 72e709b40cae · 2021-06-15T09:18:24.000-10:00
diff --git a/tests/test_aio.py b/tests/test_aio.py
@@ -18,6 +18,8 @@
 from zeroconf._services import ServiceInfo, ServiceListener
 from zeroconf._utils.time import current_time_millis
 
+from . import _clear_cache
+
 
 from . import _clear_cache
 
@@ -498,3 +500,61 @@ async def test_async_context_manager() -> None:
         await task
         aiosinfo = await aiozc.async_get_service_info(type_, registration_name)
         assert aiosinfo is not None
+
+
+@pytest.mark.asyncio
+async def test_async_unregister_all_services() -> None:
+    """Test unregistering all services."""
+    aiozc = AsyncZeroconf(interfaces=['127.0.0.1'])
+    type_ = "_test1-srvc-type._tcp.local."
+    name = "xxxyyy"
+    name2 = "abc"
+    registration_name = "%s.%s" % (name, type_)
+    registration_name2 = "%s.%s" % (name2, type_)
+
+    desc = {'path': '/~paulsm/'}
+    info = ServiceInfo(
+        type_,
+        registration_name,
+        80,
+        0,
+        0,
+        desc,
+        "ash-1.local.",
+        addresses=[socket.inet_aton("10.0.1.2")],
+    )
+    info2 = ServiceInfo(
+        type_,
+        registration_name2,
+        80,
+        0,
+        0,
+        desc,
+        "ash-5.local.",
+        addresses=[socket.inet_aton("10.0.1.5")],
+    )
+    tasks = []
+    tasks.append(await aiozc.async_register_service(info))
+    tasks.append(await aiozc.async_register_service(info2))
+    await asyncio.gather(*tasks)
+
+    tasks = []
+    tasks.append(aiozc.async_get_service_info(type_, registration_name))
+    tasks.append(aiozc.async_get_service_info(type_, registration_name2))
+    results = await asyncio.gather(*tasks)
+    assert results[0] is not None
+    assert results[1] is not None
+
+    await aiozc.async_unregister_all_services()
+
+    tasks = []
+    tasks.append(aiozc.async_get_service_info(type_, registration_name))
+    tasks.append(aiozc.async_get_service_info(type_, registration_name2))
+    results = await asyncio.gather(*tasks)
+    assert results[0] is None
+    assert results[1] is None
+
+    # Verify we can call again
+    await aiozc.async_unregister_all_services()
+
+    await aiozc.async_close()
diff --git a/zeroconf/aio.py b/zeroconf/aio.py
@@ -227,6 +227,21 @@ async def async_register_service(
         self.zeroconf.registry.add(info)
         return asyncio.ensure_future(self._async_broadcast_service(info, _REGISTER_TIME, None))
 
+    async def async_unregister_all_services(self) -> None:
+        """Unregister all registered services.
+
+        Unlike async_register_service and async_unregister_service, this
+        method does not return a future and is always expected to be
+        awaited since its only called at shutdown.
+        """
+        out = self.zeroconf.generate_unregister_all_services()
+        if not out:
+            return
+        for i in range(3):
+            if i != 0:
+                await asyncio.sleep(millis_to_seconds(_UNREGISTER_TIME))
+            self.zeroconf.async_send(out)
+
     async def async_check_service(self, info: ServiceInfo, cooperating_responders: bool = False) -> None:
         """Checks the network for a unique service name."""
         instance_name_from_service_info(info)
