Test and fix force_remove

Author: puddly

tests/test_application.py

@@ -69,6 +69,7 @@ async def callback():
 
             for response in responses:
                 await asyncio.sleep(0.1)
+                LOGGER.debug("Replying to %s with %s", request, response)
                 self.send(response)
 
             called_future.set_result(True)
@@ -84,6 +85,7 @@ async def callback():
 
             for response in responses:
                 await asyncio.sleep(0.1)
+                LOGGER.debug("Replying to %s with %s", request, response)
                 self.send(response)
 
         callback.call_count = 0
@@ -822,3 +824,47 @@ async def test_update_network_bad_channel(mocker, caplog, application):
         await app.update_network(
             channel=t.uint8_t(12), channels=t.Channels.from_channel_list([11, 15, 20]),
         )
+
+
+@pytest_mark_asyncio_timeout(seconds=3)
+async def test_force_remove(application, mocker):
+    app, znp_server = application
+
+    await app.startup(auto_form=False)
+
+    mocker.patch("zigpy_znp.zigbee.application.ZDO_REQUEST_TIMEOUT", new=0.3)
+
+    device = app.add_device(ieee=t.EUI64(range(8)), nwk=0xAABB)
+    device.status = zigpy.device.Status.ENDPOINTS_INIT
+    device.initializing = False
+
+    # Reply to zigpy's leave request
+    bad_mgmt_leave_req = znp_server.reply_once_to(
+        request=c.ZDOCommands.MgmtLeaveReq.Req(DstAddr=device.nwk, partial=True),
+        responses=[c.ZDOCommands.MgmtLeaveReq.Rsp(Status=t.Status.Failure)],
+    )
+
+    # Reply to our own leave request
+    good_mgmt_leave_req = znp_server.reply_once_to(
+        request=c.ZDOCommands.MgmtLeaveReq.Req(DstAddr=0x0000, partial=True),
+        responses=[
+            c.ZDOCommands.MgmtLeaveReq.Rsp(Status=t.Status.Success),
+            c.ZDOCommands.LeaveInd.Callback(
+                NWK=device.nwk,
+                IEEE=device.ieee,
+                Remove=False,
+                Request=False,
+                Rejoin=False,
+            ),
+        ],
+    )
+
+    # Make sure the device exists
+    assert app.get_device(nwk=device.nwk) is device
+
+    await app.remove(device.ieee)
+    await asyncio.gather(bad_mgmt_leave_req, good_mgmt_leave_req)
+
+    # Make sure the device is gone once we remove it
+    with pytest.raises(KeyError):
+        app.get_device(nwk=device.nwk)

zigpy_znp/api.py

@@ -315,7 +315,7 @@ async def request(self, request, **response_params):
         # If the sync response we got is not what we wanted, this is an error
         if not partial_response.matches(response):
             raise InvalidCommandResponse(
-                f"SRSP was not what we expected: {response} !~ {partial_response}"
+                f"Expected SRSP response {partial_response}, got {response}", response
             )
 
         return response

zigpy_znp/exceptions.py

@@ -11,4 +11,6 @@ class CommandNotRecognized(CommandError):
 
 
 class InvalidCommandResponse(CommandError):
-    pass
+    def __init__(self, message, response):
+        super().__init__(message)
+        self.response = response

zigpy_znp/zigbee/application.py

@@ -16,10 +16,12 @@
 
 from zigpy.types import ExtendedPanId, deserialize as list_deserialize
 from zigpy.zcl.clusters.security import IasZone
+from zigpy.exceptions import DeliveryError
 
 import zigpy_znp.config as conf
 import zigpy_znp.types as t
 import zigpy_znp.commands as c
+from zigpy_znp.exceptions import InvalidCommandResponse
 
 from zigpy_znp.api import ZNP
 from zigpy_znp.types.nvids import NwkNvIds
@@ -33,46 +35,45 @@
 
 ZDO_CONVERTERS = {
     ZDOCmd.Node_Desc_req: (
-        ZDOCmd.Node_Desc_rsp,
         (
-            lambda addr, NWKAddrOfInterest: c.ZDOCommands.NodeDescReq.Req(
+            lambda addr, device, NWKAddrOfInterest: c.ZDOCommands.NodeDescReq.Req(
                 DstAddr=addr, NWKAddrOfInterest=NWKAddrOfInterest
             )
         ),
-        (
-            lambda addr: c.ZDOCommands.NodeDescRsp.Callback(
-                partial=True, Src=addr, Status=t.ZDOStatus.SUCCESS
-            )
-        ),
-        (lambda rsp, dev: [rsp.NodeDescriptor]),
+        (lambda addr: c.ZDOCommands.NodeDescRsp.Callback(partial=True, Src=addr)),
+        (lambda rsp: (ZDOCmd.Node_Desc_rsp, [rsp.NodeDescriptor])),
     ),
     ZDOCmd.Active_EP_req: (
-        ZDOCmd.Active_EP_rsp,
         (
-            lambda addr, NWKAddrOfInterest: c.ZDOCommands.ActiveEpReq.Req(
+            lambda addr, device, NWKAddrOfInterest: c.ZDOCommands.ActiveEpReq.Req(
                 DstAddr=addr, NWKAddrOfInterest=NWKAddrOfInterest
             )
         ),
-        (
-            lambda addr: c.ZDOCommands.ActiveEpRsp.Callback(
-                partial=True, Src=addr, Status=t.ZDOStatus.SUCCESS
-            )
-        ),
-        (lambda rsp, dev: [rsp.ActiveEndpoints]),
+        (lambda addr: c.ZDOCommands.ActiveEpRsp.Callback(partial=True, Src=addr)),
+        (lambda rsp: (ZDOCmd.Active_EP_rsp, [rsp.ActiveEndpoints])),
     ),
     ZDOCmd.Simple_Desc_req: (
-        ZDOCmd.Simple_Desc_rsp,
         (
-            lambda addr, NWKAddrOfInterest, EndPoint: c.ZDOCommands.SimpleDescReq.Req(
+            # fmt: off
+            lambda addr, device, NWKAddrOfInterest, EndPoint: \
+            c.ZDOCommands.SimpleDescReq.Req(
                 DstAddr=addr, NWKAddrOfInterest=NWKAddrOfInterest, Endpoint=EndPoint
             )
+            # fmt: on
         ),
+        (lambda addr: c.ZDOCommands.SimpleDescRsp.Callback(partial=True, Src=addr)),
+        (lambda rsp: (ZDOCmd.Simple_Desc_rsp, [rsp.SimpleDescriptor])),
+    ),
+    ZDOCmd.Mgmt_Leave_req: (
         (
-            lambda addr: c.ZDOCommands.SimpleDescRsp.Callback(
-                partial=True, Src=addr, Status=t.ZDOStatus.SUCCESS
+            lambda addr, device, DeviceAddress, Options: c.ZDOCommands.MgmtLeaveReq.Req(
+                DstAddr=addr,
+                IEEE=device.ieee,
+                RemoveChildren_Rejoin=c.zdo.LeaveOptions(Options),
             )
         ),
-        (lambda rsp, dev: [rsp.SimpleDescriptor]),
+        (lambda addr: c.ZDOCommands.MgmtLeaveRsp.Callback(partial=True, Src=addr)),
+        (lambda rsp: (ZDOCmd.Mgmt_Leave_rsp, [rsp.Status])),
     ),
 }
 
@@ -480,34 +481,38 @@ async def _send_zdo_request(
         zdo_args, _ = list_deserialize(data, field_types)
         zdo_kwargs = dict(zip(field_names, zdo_args))
 
+        device = self.get_device(nwk=dst_addr.address)
+
         # Call the converter with the ZDO request's kwargs
-        rsp_cluster, req_factory, callback_factory, converter = ZDO_CONVERTERS[cluster]
-        request = req_factory(dst_addr.address, **zdo_kwargs)
-        callback = callback_factory(dst_addr.address)
+        req_factory, rsp_factory, zdo_rsp_factory = ZDO_CONVERTERS[cluster]
+        request = req_factory(dst_addr.address, device, **zdo_kwargs)
+        callback = rsp_factory(dst_addr.address)
 
         LOGGER.debug(
-            "Intercepted AP ZDO request and replaced with %s - %s", request, callback
+            "Intercepted AP ZDO request and replaced with %s/%s", request, callback
         )
 
-        async with async_timeout.timeout(ZDO_REQUEST_TIMEOUT):
-            response = await self._znp.request_callback_rsp(
-                request=request, RspStatus=t.Status.Success, callback=callback
-            )
+        try:
+            async with async_timeout.timeout(ZDO_REQUEST_TIMEOUT):
+                response = await self._znp.request_callback_rsp(
+                    request=request, RspStatus=t.Status.Success, callback=callback
+                )
+        except InvalidCommandResponse as e:
+            raise DeliveryError(f"Could not send command: {e.response.Status}") from e
 
-        device = self.get_device(nwk=dst_addr.address)
+        zdo_rsp_cluster, zdo_response_args = zdo_rsp_factory(response)
 
         # Build up a ZDO response
         message = t.serialize_list(
-            [t.uint8_t(sequence), response.Status, response.NWK]
-            + converter(response, device)
+            [t.uint8_t(sequence), response.Status, response.NWK] + zdo_response_args
         )
         LOGGER.trace("Pretending we received a ZDO message: %s", message)
 
         # We do not get any LQI info here
         self.handle_message(
             sender=device,
             profile=zigpy.profiles.zha.PROFILE_ID,
-            cluster=rsp_cluster,
+            cluster=zdo_rsp_cluster,
             src_ep=dst_ep,
             dst_ep=src_ep,
             message=message,
@@ -631,13 +636,22 @@ async def force_remove(self, device) -> None:
         """Forcibly remove device from NCP."""
         await self._znp.request(
             c.ZDOCommands.MgmtLeaveReq.Req(
-                DstAddr=device.nwk,
+                DstAddr=0x0000,  # We handle it
                 IEEE=device.ieee,
-                LeaveOptions=c.zdo.LeaveOptions.NONE,
+                RemoveChildren_Rejoin=c.zdo.LeaveOptions.NONE,
             ),
             RspStatus=t.Status.Success,
         )
 
+        # TODO: see what happens when we forcibly remove a device that isn't our child
+
+        # Just wait for the response, removing the device will be handled upstream
+        await self._znp.wait_for_response(
+            c.ZDOCommands.LeaveInd.Callback(
+                NWK=device.nwk, IEEE=device.ieee, partial=True
+            )
+        )
+
     async def permit_ncp(self, time_s: int) -> None:
         response = await self._znp.request_callback_rsp(
             request=c.ZDOCommands.MgmtPermitJoinReq.Req(