Improve test coverage

Author: puddly

tests/application/test_requests.py

@@ -1,7 +1,7 @@
 import asyncio
+import logging
 
 import pytest
-import zigpy.zdo
 import zigpy.endpoint
 import zigpy.profiles
 import zigpy.zdo.types as zdo_t
@@ -967,3 +967,36 @@ async def test_route_discovery_concurrency(device, make_application):
     assert route_discovery2.call_count == 2
 
     await app.shutdown()
+
+
+@pytest.mark.parametrize("device", [FormedLaunchpadCC26X2R1])
+async def test_zdo_from_unknown(device, make_application, caplog, mocker):
+    mocker.patch("zigpy_znp.zigbee.application.IEEE_ADDR_DISCOVERY_TIMEOUT", new=0.1)
+
+    app, znp_server = make_application(server_cls=device)
+
+    znp_server.reply_once_to(
+        request=c.ZDO.IEEEAddrReq.Req(partial=True),
+        responses=[c.ZDO.IEEEAddrReq.Rsp(Status=t.Status.SUCCESS)],
+    )
+
+    await app.startup(auto_form=False)
+
+    caplog.set_level(logging.WARNING)
+
+    znp_server.send(
+        c.ZDO.MsgCbIncoming.Callback(
+            Src=0x1234,
+            IsBroadcast=t.Bool.false,
+            ClusterId=zdo_t.ZDOCmd.Mgmt_Leave_rsp,
+            SecurityUse=0,
+            TSN=123,
+            MacDst=0x0000,
+            Data=t.Bytes([123, 0x00]),
+        )
+    )
+
+    await asyncio.sleep(0.5)
+    assert "unknown device" in caplog.text
+
+    await app.shutdown()

zigpy_znp/zigbee/application.py

@@ -693,11 +693,13 @@ async def on_zdo_message(self, msg: c.ZDO.MsgCbIncoming.Callback) -> None:
             self.on_zdo_device_announce(*args)
             device = self.get_device(ieee=kwargs["IEEEAddr"])
         else:
-            device = await self._get_or_discover_device(nwk=msg.Src)
-
-        if device is None:
-            LOGGER.warning("Received a ZDO message from an unknown device: %s", msg.Src)
-            return
+            try:
+                device = await self._get_or_discover_device(nwk=msg.Src)
+            except KeyError:
+                LOGGER.warning(
+                    "Received a ZDO message from an unknown device: %s", msg.Src
+                )
+                return
 
         self.handle_message(
             sender=device,