Implement network scanning using the standard zigpy interface (#648)

* Implement existing scan commands on top of `callback_for_commands`

* Implement `network_scan`

* `duration` -> `duration_exp`

* Use underscored method

* Add a unit test

Author: puddly

bellows/ezsp/__init__.py

@@ -223,28 +223,44 @@ async def _list_command(
         self, name, item_frames, completion_frame, spos, *args, **kwargs
     ):
         """Run a command, returning result callbacks as a list"""
-        fut = asyncio.Future()
+        queue = asyncio.Queue()
         results = []
 
-        def cb(frame_name, response):
-            if frame_name in item_frames:
+        with self.callback_for_commands(
+            commands=set(item_frames) | {completion_frame},
+            callback=lambda command, response: queue.put_nowait((command, response)),
+        ):
+            v = await self._command(name, *args, **kwargs)
+            if t.sl_Status.from_ember_status(v[0]) != t.sl_Status.OK:
+                raise Exception(v)
+
+            while True:
+                command, response = await queue.get()
+                if command == completion_frame:
+                    if t.sl_Status.from_ember_status(response[spos]) != t.sl_Status.OK:
+                        raise Exception(response)
+
+                    break
+
                 results.append(response)
-            elif frame_name == completion_frame:
-                fut.set_result(response)
+
+        return results
+
+    @contextlib.contextmanager
+    def callback_for_commands(
+        self, commands: set[str], callback: Callable
+    ) -> Generator[None]:
+        def cb(frame_name, response):
+            if frame_name in commands:
+                callback(frame_name, response)
 
         cbid = self.add_callback(cb)
+
         try:
-            v = await self._command(name, *args, **kwargs)
-            if t.sl_Status.from_ember_status(v[0]) != t.sl_Status.OK:
-                raise Exception(v)
-            v = await fut
-            if t.sl_Status.from_ember_status(v[spos]) != t.sl_Status.OK:
-                raise Exception(v)
+            yield
         finally:
             self.remove_callback(cbid)
 
-        return results
-
     startScan = functools.partialmethod(
         _list_command,
         "startScan",
@@ -253,7 +269,11 @@ def cb(frame_name, response):
         1,
     )
     pollForData = functools.partialmethod(
-        _list_command, "pollForData", ["pollHandler"], "pollCompleteHandler", 0
+        _list_command,
+        "pollForData",
+        ["pollHandler"],
+        "pollCompleteHandler",
+        0,
     )
     zllStartScan = functools.partialmethod(
         _list_command,

bellows/zigbee/application.py

@@ -5,6 +5,7 @@
 import os
 import statistics
 import sys
+from typing import AsyncGenerator
 
 if sys.version_info[:2] < (3, 11):
     from async_timeout import timeout as asyncio_timeout  # pragma: no cover
@@ -711,6 +712,46 @@ async def energy_scan(
             for channel in list(channels)
         }
 
+    async def _network_scan(
+        self, channels: t.Channels, duration_exp: int
+    ) -> AsyncGenerator[zigpy.types.NetworkBeacon]:
+        """Scans for networks and yields network beacons."""
+        queue = asyncio.Queue()
+
+        with self._ezsp.callback_for_commands(
+            {"networkFoundHandler", "scanCompleteHandler"},
+            callback=lambda command, response: queue.put_nowait((command, response)),
+        ):
+            # XXX: replace with normal command invocation once overload is removed
+            (status,) = await self._ezsp._command(
+                "startScan",
+                scanType=t.EzspNetworkScanType.ACTIVE_SCAN,
+                channelMask=channels,
+                duration=duration_exp,
+            )
+
+            if t.sl_Status.from_ember_status(status) != t.sl_Status.OK:
+                raise ControllerError(f"Failed to start scan: {status!r}")
+
+            while True:
+                command, response = await queue.get()
+
+                if command == "scanCompleteHandler":
+                    break
+
+                (networkFound, lastHopLqi, lastHopRssi) = response
+
+                yield zigpy.types.NetworkBeacon(
+                    pan_id=networkFound.panId,
+                    extended_pan_id=networkFound.extendedPanId,
+                    channel=networkFound.channel,
+                    nwk_update_id=networkFound.nwkUpdateId,
+                    permit_joining=bool(networkFound.allowingJoin),
+                    stack_profile=networkFound.stackProfile,
+                    lqi=lastHopLqi,
+                    rssi=lastHopRssi,
+                )
+
     async def send_packet(self, packet: zigpy.types.ZigbeePacket) -> None:
         if not self.is_controller_running:
             raise ControllerError("ApplicationController is not running")

tests/test_application.py

@@ -1943,3 +1943,104 @@ async def test_write_network_info(
             )
         )
     ]
+
+
+async def test_network_scan(app: ControllerApplication) -> None:
+    app._ezsp._protocol.startScan.return_value = [t.sl_Status.OK]
+
+    def run_scan() -> None:
+        app._ezsp._protocol._handle_callback(
+            "networkFoundHandler",
+            list(
+                {
+                    "networkFound": t.EmberZigbeeNetwork(
+                        channel=11,
+                        panId=zigpy_t.PanId(0x1D13),
+                        extendedPanId=t.EUI64.convert("00:07:81:00:00:9a:8f:3b"),
+                        allowingJoin=False,
+                        stackProfile=2,
+                        nwkUpdateId=0,
+                    ),
+                    "lastHopLqi": 152,
+                    "lastHopRssi": -62,
+                }.values()
+            ),
+        )
+        app._ezsp._protocol._handle_callback(
+            "networkFoundHandler",
+            list(
+                {
+                    "networkFound": t.EmberZigbeeNetwork(
+                        channel=11,
+                        panId=zigpy_t.PanId(0x2857),
+                        extendedPanId=t.EUI64.convert("00:07:81:00:00:9a:34:1b"),
+                        allowingJoin=False,
+                        stackProfile=2,
+                        nwkUpdateId=0,
+                    ),
+                    "lastHopLqi": 136,
+                    "lastHopRssi": -66,
+                }.values()
+            ),
+        )
+        app._ezsp._protocol._handle_callback(
+            "scanCompleteHandler",
+            list(
+                {
+                    "channel": 26,
+                    "status": t.sl_Status.OK,
+                }.values()
+            ),
+        )
+
+    asyncio.get_running_loop().call_soon(run_scan)
+
+    results = [
+        beacon
+        async for beacon in app.network_scan(
+            channels=t.Channels.from_channel_list([11, 15, 26]), duration_exp=4
+        )
+    ]
+
+    assert results == [
+        zigpy_t.NetworkBeacon(
+            pan_id=0x1D13,
+            extended_pan_id=t.EUI64.convert("00:07:81:00:00:9a:8f:3b"),
+            channel=11,
+            permit_joining=False,
+            stack_profile=2,
+            nwk_update_id=0,
+            lqi=152,
+            src=None,
+            rssi=-62,
+            depth=None,
+            router_capacity=None,
+            device_capacity=None,
+            protocol_version=None,
+        ),
+        zigpy_t.NetworkBeacon(
+            pan_id=0x2857,
+            extended_pan_id=t.EUI64.convert("00:07:81:00:00:9a:34:1b"),
+            channel=11,
+            permit_joining=False,
+            stack_profile=2,
+            nwk_update_id=0,
+            lqi=136,
+            src=None,
+            rssi=-66,
+            depth=None,
+            router_capacity=None,
+            device_capacity=None,
+            protocol_version=None,
+        ),
+    ]
+
+
+async def test_network_scan_failure(app: ControllerApplication) -> None:
+    app._ezsp._protocol.startScan.return_value = [t.sl_Status.FAIL]
+
+    with pytest.raises(zigpy.exceptions.ControllerException):
+        async for beacon in app.network_scan(
+            channels=t.Channels.from_channel_list([11, 15, 26]), duration_exp=4
+        ):
+            pass