Refactor device lookup. (#36)

Lookup devices by NWK, fallback to IEEE if it exists.

Author: Adminiuga

tests/test_application.py

@@ -76,14 +76,41 @@ def test_rx_nwk_0000(app):
     assert app._handle_reply.call_count == 0
 
 
-def test_rx_unknown_device(app):
+def test_rx_unknown_device(app, device):
+    """Unknown NWK, but existing device."""
     app._handle_reply = mock.MagicMock()
     app.handle_message = mock.MagicMock()
     app.handle_join = mock.MagicMock()
-    app.get_device = mock.MagicMock(side_effect=KeyError)
+    dev = device(nwk=0x1234)
+    app.devices[dev.ieee] = dev
+    app.get_device = mock.MagicMock(side_effect=[KeyError, dev])
+    app.deserialize = mock.MagicMock(side_effect=ValueError)
     app.handle_rx(
         b'\x01\x02\x03\x04\x05\x06\x07\x08',
-        0x1234,
+        0x3334,
+        mock.sentinel.src_ep,
+        mock.sentinel.dst_ep,
+        mock.sentinel.cluster_id,
+        mock.sentinel.profile_id,
+        mock.sentinel.rxopts,
+        b''
+    )
+    assert app.handle_join.call_count == 1
+    assert app.get_device.call_count == 2
+    assert app.handle_message.call_count == 0
+    assert app._handle_reply.call_count == 0
+
+
+def test_rx_unknown_device_iee(app):
+    """Unknown NWK, and unknown IEEE."""
+    app._handle_reply = mock.MagicMock()
+    app.handle_message = mock.MagicMock()
+    app.handle_join = mock.MagicMock()
+    app.get_device = mock.MagicMock(side_effect=KeyError)
+    app.deserialize = mock.MagicMock(side_effect=ValueError)
+    app.handle_rx(
+        b'\xff\xff\xff\xff\xff\xff\xff\xff',
+        0x3334,
         mock.sentinel.src_ep,
         mock.sentinel.dst_ep,
         mock.sentinel.cluster_id,

zigpy_xbee/types.py

@@ -155,3 +155,7 @@ def __repr__(self):
 
     def __str__(self):
         return '0x{:04x}'.format(self)
+
+
+UNKNOWN_IEEE = EUI64([uint8_t(0xff) for i in range(0, 8)])
+UNKNOWN_NWK = NWK(0xfffe)

zigpy_xbee/uart.py

@@ -22,6 +22,7 @@ def __init__(self, api, connected_future=None):
 
     def send(self, data):
         """Send data, taking care of escaping and framing"""
+        LOGGER.debug("Sending: %s", data)
         checksum = bytes([self._checksum(data)])
         frame = self.START + self._escape(len(data).to_bytes(2, 'big') + data + checksum)
         self._transport.write(frame)

zigpy_xbee/zigbee/application.py

@@ -3,15 +3,20 @@
 import logging
 
 import zigpy.application
+import zigpy.exceptions
 import zigpy.types
 import zigpy.util
+import zigpy.zdo.types
+
+from zigpy_xbee.types import UNKNOWN_IEEE
 
 
 # how long coordinator would hold message for an end device in 10ms units
 CONF_CYCLIC_SLEEP_PERIOD = 0x0300
 # end device poll timeout = 3 * SN * SP * 10ms
 CONF_POLL_TIMEOUT = 0x029b
 
+
 LOGGER = logging.getLogger(__name__)
 
 
@@ -72,6 +77,7 @@ async def startup(self, auto_form=False):
             LOGGER.debug("Coordinator %s", 'enabled' if ce else 'disabled')
         except RuntimeError as exc:
             LOGGER.debug("sending CE command: %s", exc)
+        self.add_device(self.ieee, self.nwk)
 
     async def force_remove(self, dev):
         """Forcibly remove device from NCP."""
@@ -174,11 +180,15 @@ def handle_rx(self, src_ieee, src_nwk, src_ep, dst_ep, cluster_id, profile_id, r
             self.handle_join(nwk, ieee, 0)
 
         try:
-            device = self.get_device(ieee=ember_ieee)
+            device = self.get_device(nwk=src_nwk)
         except KeyError:
-            LOGGER.debug("Received frame from unknown device: 0x%04x/%s",
-                         src_nwk, str(ember_ieee))
-            return
+            if ember_ieee != UNKNOWN_IEEE and ember_ieee in self.devices:
+                self.handle_join(src_nwk, ember_ieee, 0)
+                device = self.get_device(ieee=ember_ieee)
+            else:
+                LOGGER.debug("Received frame from unknown device: 0x%04x/%s",
+                             src_nwk, str(ember_ieee))
+                return
 
         if device.status == zigpy.device.Status.NEW and dst_ep != 0:
             # only allow ZDO responses while initializing device