Merge pull request #122 from puddly/puddly/zdo-nwk-and-ieee-addr-req

Implement ZDO converters for `NWK_addr_req` and `IEEE_addr_req`

Author: puddly

tests/application/test_joining.py

@@ -554,6 +554,7 @@ async def test_unknown_device_discovery(device, make_application, mocker):
                 Status=t.ZDOStatus.SUCCESS,
                 IEEE=existing_ieee,
                 NWK=existing_nwk + 1,
+                NumAssoc=0,
                 Index=0,
                 Devices=[],
             ),
@@ -591,6 +592,7 @@ async def test_unknown_device_discovery(device, make_application, mocker):
                 Status=t.ZDOStatus.SUCCESS,
                 IEEE=new_ieee,
                 NWK=new_nwk,
+                NumAssoc=0,
                 Index=0,
                 Devices=[],
             ),
@@ -603,3 +605,24 @@ async def test_unknown_device_discovery(device, make_application, mocker):
     assert new_dev.ieee == new_ieee
 
     await app.pre_shutdown()
+
+
+@pytest.mark.parametrize("device", FORMED_DEVICES)
+async def test_unknown_device_discovery_failure(device, make_application, mocker):
+    mocker.patch("zigpy_znp.zigbee.application.IEEE_ADDR_DISCOVERY_TIMEOUT", new=0.1)
+
+    app, znp_server = make_application(server_cls=device)
+    await app.startup(auto_form=False)
+
+    znp_server.reply_once_to(
+        request=c.ZDO.IEEEAddrReq.Req(partial=True),
+        responses=[
+            c.ZDO.IEEEAddrReq.Rsp(Status=t.Status.SUCCESS),
+        ],
+    )
+
+    # Discovery will throw an exception when the device cannot be found
+    with pytest.raises(KeyError):
+        await app._get_or_discover_device(nwk=0x3456)
+
+    await app.pre_shutdown()

tests/application/test_zigpy_callbacks.py

@@ -10,15 +10,20 @@
 from ..conftest import FORMED_DEVICES, CoroutineMock
 
 
-def awaitable_mock(return_value):
+def awaitable_mock(*, return_value=None, side_effect=None):
+    assert (return_value or side_effect) and not (return_value and side_effect)
+
     mock_called = asyncio.get_running_loop().create_future()
 
-    def side_effect(*args, **kwargs):
+    def side_effect_(*args, **kwargs):
         mock_called.set_result((args, kwargs))
 
-        return return_value
+        if return_value is not None:
+            return return_value
+        else:
+            raise side_effect
 
-    return mock_called, CoroutineMock(side_effect=side_effect)
+    return mock_called, CoroutineMock(side_effect=side_effect_)
 
 
 @pytest.mark.parametrize("device", FORMED_DEVICES)
@@ -45,7 +50,7 @@ async def test_on_zdo_relays_message_callback_unknown(
     app, znp_server = make_application(server_cls=device)
     await app.startup(auto_form=False)
 
-    discover_called, discover_mock = awaitable_mock(return_value=None)
+    discover_called, discover_mock = awaitable_mock(side_effect=KeyError())
     mocker.patch.object(app, "_get_or_discover_device", new=discover_mock)
 
     caplog.set_level(logging.WARNING)
@@ -183,7 +188,7 @@ async def test_on_af_message_callback(device, make_application, mocker):
     app.get_device.reset_mock()
 
     # Message from an unknown device
-    discover_called, discover_mock = awaitable_mock(return_value=None)
+    discover_called, discover_mock = awaitable_mock(side_effect=KeyError())
     mocker.patch.object(app, "_get_or_discover_device", new=discover_mock)
 
     znp_server.send(af_message)

zigpy_znp/commands/zdo.py

@@ -116,6 +116,10 @@ class ChildInfoList(t.LVList, item_type=t.EUI64, length_type=t.uint8_t):
     pass
 
 
+class NWKArray(t.CompleteList, item_type=t.NWK):
+    pass
+
+
 class NullableNodeDescriptor(zigpy.zdo.types.NodeDescriptor):
     @classmethod
     def deserialize(cls, data: bytes) -> tuple[NullableNodeDescriptor, bytes]:
@@ -935,12 +939,13 @@ class ZDO(t.CommandsBase, subsystem=t.Subsystem.ZDO):
             ),
             t.Param("IEEE", t.EUI64, "Extended address of the source device"),
             t.Param("NWK", t.NWK, "Short address of the source device"),
+            t.Param("NumAssoc", t.uint8_t, "Number of associated devices"),
             t.Param(
                 "Index",
                 t.uint8_t,
                 "Starting index into the list of associated devices",
             ),
-            t.Param("Devices", t.NWKList, "List of the associated devices"),
+            t.Param("Devices", NWKArray, "List of the associated devices"),
         ),
     )
 
@@ -954,12 +959,13 @@ class ZDO(t.CommandsBase, subsystem=t.Subsystem.ZDO):
             ),
             t.Param("IEEE", t.EUI64, "Extended address of the source device"),
             t.Param("NWK", t.NWK, "Short address of the source device"),
+            t.Param("NumAssoc", t.uint8_t, "Number of associated devices"),
             t.Param(
                 "Index",
                 t.uint8_t,
                 "Starting index into the list of associated devices",
             ),
-            t.Param("Devices", t.NWKList, "List of the associated devices"),
+            t.Param("Devices", NWKArray, "List of the associated devices"),
         ),
     )
 

zigpy_znp/utils.py

@@ -164,12 +164,14 @@ def matches(self, other) -> bool:
         return True
 
 
-def combine_concurrent_calls(function):
+def combine_concurrent_calls(
+    function: typing.CoroutineFunction,
+) -> typing.CoroutineFunction:
     """
     Decorator that allows concurrent calls to expensive coroutines to share a result.
     """
 
-    futures = {}
+    tasks = {}
     signature = inspect.signature(function)
 
     @functools.wraps(function)
@@ -180,20 +182,15 @@ async def replacement(*args, **kwargs):
         # XXX: all args and kwargs are assumed to be hashable
         key = tuple(bound.arguments.items())
 
-        if key in futures:
-            return await futures[key]
+        if key in tasks:
+            return await tasks[key]
 
-        future = futures[key] = asyncio.get_running_loop().create_future()
+        tasks[key] = asyncio.create_task(function(*args, **kwargs))
 
         try:
-            result = await function(*args, **kwargs)
-        except Exception as e:
-            future.set_exception(e)
-            raise
-        else:
-            future.set_result(result)
-            return result
+            return await tasks[key]
         finally:
-            del futures[key]
+            assert tasks[key].done()
+            del tasks[key]
 
     return replacement

zigpy_znp/zigbee/application.py

@@ -47,6 +47,7 @@
 STARTUP_TIMEOUT = 5
 ZDO_REQUEST_TIMEOUT = 15
 DATA_CONFIRM_TIMEOUT = 8
+IEEE_ADDR_DISCOVERY_TIMEOUT = 5
 DEVICE_JOIN_MAX_DELAY = 5
 WATCHDOG_PERIOD = 30
 BROADCAST_SEND_WAIT_DURATION = 3
@@ -670,6 +671,12 @@ def _bind_callbacks(self) -> None:
             self.on_intentionally_unhandled_message,
         )
 
+        # These are responses to a broadcast but we ignore all but the first
+        self._znp.callback_for_response(
+            c.ZDO.IEEEAddrRsp.Callback(partial=True),
+            self.on_intentionally_unhandled_message,
+        )
+
     def on_intentionally_unhandled_message(self, msg: t.CommandBase) -> None:
         """
         Some commands are unhandled but frequently sent by devices on the network. To
@@ -693,9 +700,9 @@ async def on_zdo_relays_message(self, msg: c.ZDO.SrcRtgInd.Callback) -> None:
         ZDO source routing message callback
         """
 
-        device = await self._get_or_discover_device(nwk=msg.DstAddr)
-
-        if device is None:
+        try:
+            device = await self._get_or_discover_device(nwk=msg.DstAddr)
+        except KeyError:
             LOGGER.warning(
                 "Received a ZDO message from an unknown device: %s", msg.DstAddr
             )
@@ -778,9 +785,9 @@ async def on_af_message(self, msg: c.AF.IncomingMsg.Callback) -> None:
         Handler for all non-ZDO messages.
         """
 
-        device = await self._get_or_discover_device(nwk=msg.SrcAddr)
-
-        if device is None:
+        try:
+            device = await self._get_or_discover_device(nwk=msg.SrcAddr)
+        except KeyError:
             LOGGER.warning(
                 "Received an AF message from an unknown device: %s", msg.SrcAddr
             )
@@ -864,7 +871,7 @@ async def _watchdog_loop(self):
                 return
 
     @combine_concurrent_calls
-    async def _get_or_discover_device(self, nwk: t.NWK) -> zigpy.device.Device | None:
+    async def _get_or_discover_device(self, nwk: t.NWK) -> zigpy.device.Device:
         """
         Finds a device by its NWK address. If a device does not exist in the zigpy
         database, attempt to look up its new NWK address. If it does not exist in the
@@ -880,23 +887,16 @@ async def _get_or_discover_device(self, nwk: t.NWK) -> zigpy.device.Device | Non
 
         try:
             # XXX: Multiple responses may arrive but we only use the first one
-            ieee_addr_rsp = await self._znp.request_callback_rsp(
-                request=c.ZDO.IEEEAddrReq.Req(
-                    NWK=nwk,
-                    RequestType=c.zdo.AddrRequestType.SINGLE,
-                    StartIndex=0,
-                ),
-                RspStatus=t.Status.SUCCESS,
-                callback=c.ZDO.IEEEAddrRsp.Callback(
-                    partial=True,
-                    NWK=nwk,
-                ),
-                timeout=5,  # We don't want to wait forever
-            )
+            async with async_timeout.timeout(IEEE_ADDR_DISCOVERY_TIMEOUT):
+                _, ieee, _, _, _, _ = await self.zigpy_device.zdo.IEEE_addr_req(
+                    *{
+                        "NWKAddrOfInterest": nwk,
+                        "RequestType": c.zdo.AddrRequestType.SINGLE,
+                        "StartIndex": 0,
+                    }.values()
+                )
         except asyncio.TimeoutError:
-            return None
-
-        ieee = ieee_addr_rsp.IEEE
+            raise KeyError(f"Unknown device: 0x{nwk:04X}")
 
         try:
             device = self.get_device(ieee=ieee)
@@ -1276,7 +1276,7 @@ async def _send_zdo_request(
         # Call the converter with the ZDO request's kwargs
         req_factory, rsp_factory, zdo_rsp_factory = ZDO_CONVERTERS[cluster]
         request = req_factory(dst_addr, **zdo_kwargs)
-        callback = rsp_factory(dst_addr)
+        callback = rsp_factory(dst_addr, **zdo_kwargs)
 
         LOGGER.debug(
             "Intercepted AP ZDO request %s(%s) and replaced with %s",