feat: context managers in ServiceBrowser and AsyncServiceBrowser

Tenebrosus3141 · web-flow · commit cb069efdb8db · 2023-08-27T01:23:29.000Z
diff --git a/src/zeroconf/_services/browser.py b/src/zeroconf/_services/browser.py
@@ -26,6 +26,7 @@
 import threading
 import warnings
 from abc import abstractmethod
+from types import TracebackType  # noqa # used in type hints
 from typing import (
     TYPE_CHECKING,
     Callable,
@@ -35,6 +36,7 @@
     Optional,
     Set,
     Tuple,
+    Type,
     Union,
     cast,
 )
@@ -576,3 +578,15 @@ def async_update_records_complete(self) -> None:
         for pending in self._pending_handlers.items():
             self.queue.put(pending)
         self._pending_handlers.clear()
+
+    def __enter__(self) -> 'ServiceBrowser':
+        return self
+
+    def __exit__(  # pylint: disable=useless-return
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> Optional[bool]:
+        self.cancel()
+        return None
diff --git a/src/zeroconf/asyncio.py b/src/zeroconf/asyncio.py
@@ -93,6 +93,18 @@ def async_update_records_complete(self) -> None:
             self._fire_service_state_changed_event(pending)
         self._pending_handlers.clear()
 
+    async def __aenter__(self) -> 'AsyncServiceBrowser':
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> Optional[bool]:
+        await self.async_cancel()
+        return None
+
 
 class AsyncZeroconfServiceTypes(ZeroconfServiceTypes):
     """An async version of ZeroconfServiceTypes."""
diff --git a/tests/services/test_browser.py b/tests/services/test_browser.py
@@ -9,7 +9,7 @@
 import time
 import unittest
 from threading import Event
-from typing import Iterable, Set
+from typing import cast, Iterable, Set
 from unittest.mock import patch
 
 import pytest
@@ -70,6 +70,31 @@ class MyServiceListener(r.ServiceListener):
     zc.close()
 
 
+def test_service_browser_cancel_context_manager():
+    """Test we can cancel a ServiceBrowser with it being used as a context manager."""
+
+    # instantiate a zeroconf instance
+    zc = Zeroconf(interfaces=['127.0.0.1'])
+    # start a browser
+    type_ = "_hap._tcp.local."
+
+    class MyServiceListener(r.ServiceListener):
+        pass
+
+    listener = MyServiceListener()
+
+    browser = r.ServiceBrowser(zc, type_, None, listener)
+
+    assert cast(browser.done, bool) is False
+
+    with browser:
+        pass
+
+    assert cast(browser.done, bool) is True
+
+    zc.close()
+
+
 def test_service_browser_cancel_multiple_times_after_close():
     """Test we can cancel a ServiceBrowser multiple times after close."""
 
diff --git a/tests/test_asyncio.py b/tests/test_asyncio.py
@@ -9,6 +9,7 @@
 import socket
 import threading
 import time
+from typing import cast
 from unittest.mock import ANY, call, patch
 
 import pytest
@@ -779,6 +780,34 @@ async def test_async_context_manager() -> None:
         assert aiosinfo is not None
 
 
+@pytest.mark.asyncio
+async def test_service_browser_cancel_async_context_manager():
+    """Test we can cancel an AsyncServiceBrowser with it being used as an async context manager."""
+
+    # instantiate a zeroconf instance
+    aiozc = AsyncZeroconf(interfaces=['127.0.0.1'])
+    zc = aiozc.zeroconf
+    type_ = "_hap._tcp.local."
+
+    class MyServiceListener(ServiceListener):
+        pass
+
+    listener = MyServiceListener()
+
+    browser = AsyncServiceBrowser(zc, type_, None, listener)
+
+    assert cast(browser.done, bool) is False
+
+    async with browser:
+        pass
+
+    await asyncio.sleep(0)
+
+    assert cast(browser.done, bool) is True
+
+    await aiozc.async_close()
+
+
 @pytest.mark.asyncio
 async def test_async_unregister_all_services() -> None:
     """Test unregistering all services."""
