Test energy scanner tool

Author: puddly

tests/test_application.py

@@ -162,8 +162,7 @@ def passthrough_serial_conn(loop, protocol_factory, url, *args, **kwargs):
     return server_znp
 
 
-@pytest.fixture
-def application(znp_server):
+def make_application(znp_server):
     app = ControllerApplication(config_for_port_path("/dev/ttyFAKE0"))
 
     # Handle the entire startup sequence
@@ -291,6 +290,11 @@ def on_endpoint_deletion(req):
     return app, znp_server
 
 
+@pytest.fixture
+def application(znp_server):
+    return make_application(znp_server)
+
+
 @pytest_mark_asyncio_timeout(seconds=5)
 async def test_application_startup_skip_bootloader(application, mocker):
     app, znp_server = application

tests/test_tools_energy_scan.py

@@ -0,0 +1,64 @@
+import asyncio
+
+import zigpy_znp.types as t
+import zigpy_znp.commands as c
+
+from zigpy_znp.tools.energy_scan import channels_from_channel_mask, main as energy_scan
+
+from test_api import pytest_mark_asyncio_timeout  # noqa: F401
+from test_application import make_application, znp_server  # noqa: F401
+from test_tools_nvram import openable_serial_znp_server  # noqa: F401
+
+
+def test_channels_from_channel_mask():
+    def channel_list(channels):
+        return list(channels_from_channel_mask(channels))
+
+    assert channel_list(t.Channels.ALL_CHANNELS) == list(range(11, 26 + 1))
+    assert channel_list(t.Channels.NO_CHANNELS) == []
+    assert channel_list(t.Channels.CHANNEL_11 | t.Channels.CHANNEL_20) == [11, 20]
+    assert channel_list(t.Channels.CHANNEL_15) == [15]
+
+
+@pytest_mark_asyncio_timeout(seconds=5)
+async def test_energy_scan(openable_serial_znp_server, capsys):  # noqa: F811
+    app, openable_serial_znp_server = make_application(openable_serial_znp_server)
+
+    def fake_scanner(request):
+        async def response(request):
+            openable_serial_znp_server.send(
+                c.ZDO.MgmtNWKUpdateReq.Rsp(Status=t.Status.SUCCESS)
+            )
+
+            delay = 2 ** request.ScanDuration
+            num_channels = len(list(channels_from_channel_mask(request.Channels)))
+
+            for i in range(request.ScanCount):
+                await asyncio.sleep(delay / 100)
+
+                openable_serial_znp_server.send(
+                    c.ZDO.MgmtNWKUpdateNotify.Callback(
+                        Src=0x0000,
+                        Status=t.ZDOStatus.SUCCESS,
+                        ScannedChannels=request.Channels,
+                        TotalTransmissions=998,
+                        TransmissionFailures=2,
+                        EnergyValues=list(range(num_channels)),
+                    )
+                )
+
+        asyncio.create_task(response(request))
+
+    openable_serial_znp_server.callback_for_response(
+        c.ZDO.MgmtNWKUpdateReq.Req(
+            Dst=0x0000, DstAddrMode=t.AddrMode.NWK, NwkManagerAddr=0x0000, partial=True,
+        ),
+        fake_scanner,
+    )
+
+    await energy_scan(["-n", "1", openable_serial_znp_server._port_path, "-v", "-v"])
+
+    captured = capsys.readouterr()
+
+    for i in range(11, 26 + 1):
+        assert str(i) in captured.out

zigpy_znp/tools/energy_scan.py

@@ -2,6 +2,7 @@
 import asyncio
 import logging
 import argparse
+import itertools
 
 from collections import defaultdict, deque
 
@@ -19,7 +20,7 @@ def channels_from_channel_mask(channels: t.Channels):
             yield channel
 
 
-async def perform_energy_scan(radio_path, auto_form=False):
+async def perform_energy_scan(radio_path, num_scans=None, auto_form=False):
     LOGGER.info("Starting up zigpy-znp")
 
     app = ControllerApplication(
@@ -30,18 +31,18 @@ async def perform_energy_scan(radio_path, auto_form=False):
         await app.startup(auto_form=auto_form, write_nvram=auto_form)
     except RuntimeError:
         LOGGER.error("The hardware needs to be configured before this tool can work.")
-        LOGGER.error(
-            "Either re-run this command with -f (--form) or use the hardware "
-            "once with either ZHA or Zigbee2Mqtt."
-        )
+        LOGGER.error("Re-run this command with -f (--form).")
         return
 
     LOGGER.info("Running scan...")
 
     # We compute an average over the last 5 scans
     channel_energies = defaultdict(lambda: deque([], maxlen=5))
 
-    while True:
+    for i in itertools.count(start=1):
+        if num_scans is not None and i > num_scans:
+            break
+
         rsp = await app._znp.request_callback_rsp(
             request=c.ZDO.MgmtNWKUpdateReq.Req(
                 Dst=0x0000,
@@ -61,7 +62,7 @@ async def perform_energy_scan(radio_path, auto_form=False):
             energies = channel_energies[channel]
             energies.append(energy)
 
-        total = 0xFF * energies.maxlen
+        total = 0xFF * len(energies)
 
         print(f"Channel energy ({len(energies)} / {energies.maxlen}):")
 
@@ -89,6 +90,14 @@ async def main(argv):
         default=0,
         help="Increases verbosity",
     )
+    parser.add_argument(
+        "-n",
+        "--num-scans",
+        dest="num_scans",
+        type=int,
+        default=None,
+        help="Number of scans to perform before exiting",
+    )
     parser.add_argument(
         "-f",
         "--form",
@@ -107,7 +116,9 @@ async def main(argv):
     # We just want to make sure it exists
     args.serial.close()
 
-    await perform_energy_scan(args.serial.name, auto_form=args.form)
+    await perform_energy_scan(
+        args.serial.name, auto_form=args.form, num_scans=args.num_scans
+    )
 
 
 if __name__ == "__main__":