From 74a2493442e8d7b885f405464e72d920efe5b1a6 Mon Sep 17 00:00:00 2001 From: puddly <32534428+puddly@users.noreply.github.com> Date: Mon, 19 Aug 2024 15:33:11 -0400 Subject: [PATCH 1/5] Implement existing scan commands on top of `callback_for_commands` --- bellows/ezsp/__init__.py | 48 ++++++++++++++++++++++++++++------------ 1 file changed, 34 insertions(+), 14 deletions(-) diff --git a/bellows/ezsp/__init__.py b/bellows/ezsp/__init__.py index c3649462..2f367812 100644 --- a/bellows/ezsp/__init__.py +++ b/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, From 1d66e45d061b22bdb051d6b64709074ee95e6a7f Mon Sep 17 00:00:00 2001 From: puddly <32534428+puddly@users.noreply.github.com> Date: Mon, 19 Aug 2024 15:33:46 -0400 Subject: [PATCH 2/5] Implement `network_scan` --- bellows/zigbee/application.py | 38 +++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py index 13b6386c..35a444e3 100644 --- a/bellows/zigbee/application.py +++ b/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,43 @@ async def energy_scan( for channel in list(channels) } + async def network_scan( + self, channels: t.Channels, duration: 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, + ) + + 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") From 84cfa35750b745b4137837cd8793aa26f1e0b5c8 Mon Sep 17 00:00:00 2001 From: puddly <32534428+puddly@users.noreply.github.com> Date: Mon, 19 Aug 2024 15:42:05 -0400 Subject: [PATCH 3/5] `duration` -> `duration_exp` --- bellows/zigbee/application.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py index 35a444e3..81f10468 100644 --- a/bellows/zigbee/application.py +++ b/bellows/zigbee/application.py @@ -713,7 +713,7 @@ async def energy_scan( } async def network_scan( - self, channels: t.Channels, duration: int + self, channels: t.Channels, duration_exp: int ) -> AsyncGenerator[zigpy.types.NetworkBeacon]: """Scans for networks and yields network beacons.""" queue = asyncio.Queue() @@ -727,7 +727,7 @@ async def network_scan( "startScan", scanType=t.EzspNetworkScanType.ACTIVE_SCAN, channelMask=channels, - duration=duration, + duration=duration_exp, ) while True: From 660bbfa67b312953b802deca0b1812172f332f5e Mon Sep 17 00:00:00 2001 From: puddly <32534428+puddly@users.noreply.github.com> Date: Thu, 16 Jan 2025 10:46:24 -0500 Subject: [PATCH 4/5] Use underscored method --- bellows/zigbee/application.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py index 81f10468..03ad9ad9 100644 --- a/bellows/zigbee/application.py +++ b/bellows/zigbee/application.py @@ -712,7 +712,7 @@ async def energy_scan( for channel in list(channels) } - async def network_scan( + async def _network_scan( self, channels: t.Channels, duration_exp: int ) -> AsyncGenerator[zigpy.types.NetworkBeacon]: """Scans for networks and yields network beacons.""" From a5bfa4b2eed5cd8c9ce63022db0c5b17b16d81a9 Mon Sep 17 00:00:00 2001 From: puddly <32534428+puddly@users.noreply.github.com> Date: Thu, 16 Jan 2025 12:03:11 -0500 Subject: [PATCH 5/5] Add a unit test --- bellows/zigbee/application.py | 3 + tests/test_application.py | 101 ++++++++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+) diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py index 03ad9ad9..068034a6 100644 --- a/bellows/zigbee/application.py +++ b/bellows/zigbee/application.py @@ -730,6 +730,9 @@ async def _network_scan( 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() diff --git a/tests/test_application.py b/tests/test_application.py index c3931c3e..873f6921 100644 --- a/tests/test_application.py +++ b/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