Add unit tests and properly handle unicast ZDO requests to coordinator

Author: puddly

tests/application/test_zdo_requests.py

@@ -0,0 +1,100 @@
+import asyncio
+
+import pytest
+import zigpy.zdo
+import zigpy.types as zigpy_t
+import zigpy.zdo.types as zdo_t
+
+import zigpy_znp.types as t
+import zigpy_znp.commands as c
+
+from tests.conftest import FormedLaunchpadCC26X2R1
+
+
+@pytest.mark.parametrize(
+    "broadcast,nwk_update_id,change_channel",
+    [
+        (False, 1, False),
+        (False, 1, True),
+        (True, 1, False),
+        (False, 200, True),
+    ],
+)
+@pytest.mark.parametrize("device", [FormedLaunchpadCC26X2R1])
+async def test_mgmt_nwk_update_req(
+    device, broadcast, nwk_update_id, change_channel, make_application, mocker
+):
+    mocker.patch("zigpy_znp.zigbee.device.NWK_UPDATE_LOOP_DELAY", 0.1)
+
+    app, znp_server = make_application(server_cls=device)
+
+    if change_channel:
+        new_channel = 11 + (26 - znp_server.nib.nwkLogicalChannel)
+    else:
+        new_channel = znp_server.nib.nwkLogicalChannel
+
+    async def update_channel(req):
+        # Wait a bit before updating
+        await asyncio.sleep(0.5)
+
+        znp_server.nib = znp_server.nib.replace(
+            nwkUpdateId=znp_server.nib.nwkUpdateId + 1,
+            nwkLogicalChannel=list(req.Channels)[0],
+            channelList=req.Channels,
+        )
+
+        yield
+
+    znp_server.reply_once_to(
+        request=c.AF.DataRequestExt.Req(
+            DstEndpoint=0,
+            ClusterId=zdo_t.ZDOCmd.Mgmt_NWK_Update_req,
+            partial=True,
+        ),
+        responses=[c.AF.DataRequestExt.Rsp(Status=t.Status.SUCCESS)],
+    )
+
+    nwk_update_req = znp_server.reply_once_to(
+        request=c.ZDO.MgmtNWKUpdateReq.Req(
+            Dst=0x0000,
+            DstAddrMode=t.AddrMode.NWK,
+            Channels=t.Channels.from_channel_list([new_channel]),
+            ScanDuration=254,
+            # Missing fields in the request cannot be `None` in the Z-Stack command
+            ScanCount=0,
+            NwkManagerAddr=0x0000,
+        ),
+        responses=[
+            c.ZDO.MgmtNWKUpdateReq.Rsp(Status=t.Status.SUCCESS),
+            update_channel,
+        ],
+    )
+
+    await app.startup(auto_form=False)
+
+    update = zdo_t.NwkUpdate(
+        ScanChannels=t.Channels.from_channel_list([new_channel]),
+        ScanDuration=zdo_t.NwkUpdate.CHANNEL_CHANGE_REQ,
+        nwkUpdateId=nwk_update_id,
+    )
+
+    if broadcast:
+        await zigpy.zdo.broadcast(
+            app,
+            zdo_t.ZDOCmd.Mgmt_NWK_Update_req,
+            0x0000,  # group id (ignore)
+            0,  # radius
+            update,
+            broadcast_address=zigpy_t.BroadcastAddress.ALL_ROUTERS_AND_COORDINATOR,
+        )
+    else:
+        await app.zigpy_device.zdo.Mgmt_NWK_Update_req(update)
+
+    if change_channel:
+        await nwk_update_req
+    else:
+        assert not nwk_update_req.done()
+
+    assert znp_server.nib.nwkLogicalChannel == list(update.ScanChannels)[0]
+
+    await app.shutdown()

zigpy_znp/zigbee/device.py

@@ -1,12 +1,16 @@
+from __future__ import annotations
+
 import asyncio
 import logging
 
 import zigpy.zdo
 import zigpy.device
 import zigpy.zdo.types as zdo_t
+import zigpy.application
 
 import zigpy_znp.types as t
 import zigpy_znp.commands as c
+import zigpy_znp.zigbee.application as znp_app
 
 LOGGER = logging.getLogger(__name__)
 
@@ -40,31 +44,61 @@ def model(self):
 
         return f"{model}, Z-Stack {version} (build {self.application._zstack_build_id})"
 
+    def request(
+        self,
+        profile,
+        cluster,
+        src_ep,
+        dst_ep,
+        sequence,
+        data,
+        expect_reply=True,
+        # Extend the default timeout
+        timeout=2 * zigpy.device.APS_REPLY_TIMEOUT,
+        use_ieee=False,
+    ):
+        """
+        Normal `zigpy.device.Device:request` except its default timeout is longer.
+        """
+
+        return super().request(
+            profile,
+            cluster,
+            src_ep,
+            dst_ep,
+            sequence,
+            data,
+            expect_reply=expect_reply,
+            timeout=timeout,
+            use_ieee=use_ieee,
+        )
+
 
 class ZNPZDOEndpoint(zigpy.zdo.ZDO):
     @property
-    def app(self):
+    def app(self) -> zigpy.application.ControllerApplication:
         return self.device.application
 
-    async def async_handle_mgmt_permit_joining_req(
-        self,
-        hdr: zdo_t.ZDOHeader,
-        PermitDuration: t.uint8_t,
-        TC_Significant: t.Bool,
-        *,
-        dst_addressing,
+    def _send_loopback_reply(
+        self, command_id: zdo_t.ZDOCmd, *, tsn: t.uint8_t, **kwargs
     ):
-        # Joins *must* be sent via a ZDO command. Otherwise, Z-Stack will not actually
-        # permit the coordinator to send the network key while routers will.
-        await self.app._znp.request_callback_rsp(
-            request=c.ZDO.MgmtPermitJoinReq.Req(
-                AddrMode=t.AddrMode.NWK,
-                Dst=0x0000,
-                Duration=PermitDuration,
-                TCSignificance=TC_Significant,
-            ),
-            RspStatus=t.Status.SUCCESS,
-            callback=c.ZDO.MgmtPermitJoinRsp.Callback(Src=0x0000, partial=True),
+        """
+        Constructs and sends back a loopback ZDO response.
+        """
+
+        message = t.uint8_t(tsn).serialize() + self._serialize(
+            command_id, *kwargs.values()
+        )
+
+        LOGGER.debug("Sending loopback reply %s (%s), tsn=%s", command_id, kwargs, tsn)
+
+        self.app.handle_message(
+            sender=self.app.zigpy_device,
+            profile=znp_app.ZDO_PROFILE,
+            cluster=command_id,
+            src_ep=znp_app.ZDO_ENDPOINT,
+            dst_ep=znp_app.ZDO_ENDPOINT,
+            message=message,
         )
 
     def handle_mgmt_nwk_update_req(
@@ -83,7 +117,7 @@ def handle_mgmt_nwk_update_req(
     async def async_handle_mgmt_nwk_update_req(
         self, hdr: zdo_t.ZDOHeader, NwkUpdate: zdo_t.NwkUpdate, *, dst_addressing
     ):
-        # Energy scans are handled properly by Z-Stack
+        # Energy scans are handled properly by Z-Stack, no need to do anything
         if NwkUpdate.ScanDuration not in (
             zdo_t.NwkUpdate.CHANNEL_CHANGE_REQ,
             zdo_t.NwkUpdate.CHANNEL_MASK_MANAGER_ADDR_CHANGE_REQ,
@@ -97,6 +131,15 @@ async def async_handle_mgmt_nwk_update_req(
             == NwkUpdate.ScanChannels
         ):
             LOGGER.warning("NWK update request is ignored when channel does not change")
+            self._send_loopback_reply(
+                zdo_t.ZDOCmd.Mgmt_NWK_Update_rsp,
+                Status=zdo_t.Status.SUCCESS,
+                ScannedChannels=t.Channels.NO_CHANNELS,
+                TotalTransmissions=0,
+                TransmissionFailures=0,
+                EnergyValues=[],
+                tsn=hdr.tsn,
+            )
             return
 
         await self.app._znp.request(
@@ -128,3 +171,13 @@ async def async_handle_mgmt_nwk_update_req(
                 f" {self.app.state.network_information.nwk_update_id} instead of being"
                 f" set to {NwkUpdate.nwkUpdateId}"
             )
+
+        self._send_loopback_reply(
+            zdo_t.ZDOCmd.Mgmt_NWK_Update_rsp,
+            Status=zdo_t.Status.SUCCESS,
+            ScannedChannels=t.Channels.NO_CHANNELS,
+            TotalTransmissions=0,
+            TransmissionFailures=0,
+            EnergyValues=[],
+            tsn=hdr.tsn,
+        )