Source routing fixes (#626)

* Enable address discovery when source routing is in use

* Only use `setSourceRoute` < EZSPv8

* Add back missing branch to guard against missing source route

* Implement `set_source_route` on the EZSP class instead of the application

* Add unit tests

* Fix unit tests

* Move `set_source_route` stub to EZSPv9, not v8

bellows/ezsp/protocol.py

@@ -239,3 +239,7 @@ async def send_broadcast(
         data: bytes,
     ) -> tuple[t.sl_Status, t.uint8_t]:
         raise NotImplementedError
+
+    @abc.abstractmethod
+    async def set_source_route(self, nwk: t.NWK, relays: list[t.NWK]) -> t.sl_Status:
+        raise NotImplementedError

bellows/ezsp/v4/__init__.py

@@ -176,3 +176,7 @@ async def send_broadcast(
         )
 
         return t.sl_Status.from_ember_status(status), sequence
+
+    async def set_source_route(self, nwk: t.NWK, relays: list[t.NWK]) -> t.sl_Status:
+        (res,) = await self.setSourceRoute(destination=nwk, relayList=relays)
+        return t.sl_Status.from_ember_status(res)

bellows/ezsp/v9/__init__.py

@@ -38,3 +38,7 @@ async def write_child_data(self, children: dict[t.EUI64, t.NWK]) -> None:
                     timeout=0,
                 ),
             )
+
+    async def set_source_route(self, nwk: t.NWK, relays: list[t.NWK]) -> t.sl_Status:
+        # While the command may succeed, it does absolutely nothing
+        return t.sl_Status.OK

bellows/zigbee/application.py

@@ -671,12 +671,6 @@ async def _reset_mfg_id(self, mfg_id: int) -> None:
         await asyncio.sleep(MFG_ID_RESET_DELAY)
         await self._ezsp.setManufacturerCode(code=DEFAULT_MFG_ID)
 
-    async def _set_source_route(
-        self, nwk: zigpy.types.NWK, relays: list[zigpy.types.NWK]
-    ) -> bool:
-        (res,) = await self._ezsp.setSourceRoute(destination=nwk, relayList=relays)
-        return t.sl_Status.from_ember_status(res) == t.sl_Status.OK
-
     async def energy_scan(
         self, channels: t.Channels, duration_exp: int, count: int
     ) -> dict[int, float]:
@@ -743,6 +737,9 @@ async def send_packet(self, packet: zigpy.types.ZigbeePacket) -> None:
 
         if not self.config[zigpy.config.CONF_SOURCE_ROUTING]:
             aps_frame.options |= t.EmberApsOption.APS_OPTION_ENABLE_ROUTE_DISCOVERY
+        else:
+            # Source routing uses address discovery to discover routes
+            aps_frame.options |= t.EmberApsOption.APS_OPTION_ENABLE_ADDRESS_DISCOVERY
 
         async with self._limit_concurrency():
             message_tag = self.get_sequence()
@@ -756,14 +753,10 @@ async def send_packet(self, packet: zigpy.types.ZigbeePacket) -> None:
                                     remoteEui64=device.ieee, extendedTimeout=True
                                 )
 
-                            if (
-                                packet.source_route is not None
-                                and not await self._set_source_route(
-                                    packet.dst.address, packet.source_route
-                                )
-                            ):
-                                aps_frame.options |= (
-                                    t.EmberApsOption.APS_OPTION_ENABLE_ROUTE_DISCOVERY
+                            if packet.source_route is not None:
+                                await self._ezsp.set_source_route(
+                                    nwk=packet.dst.address,
+                                    relays=packet.source_route,
                                 )
 
                             status, _ = await self._ezsp.send_unicast(

tests/test_application.py

@@ -831,20 +831,23 @@ async def test_send_packet_unicast_ieee_no_fallback(app, packet, caplog):
     assert app._ezsp.send_unicast.call_count == 0
 
 
-async def test_send_packet_unicast_source_route_ezsp7(make_app, packet):
+async def test_send_packet_unicast_source_route(make_app, packet):
     app = make_app({zigpy.config.CONF_SOURCE_ROUTING: True})
-    app._ezsp.ezsp_version = 7
-    app._ezsp.setSourceRoute = AsyncMock(return_value=(t.EmberStatus.SUCCESS,))
+    app._ezsp.set_source_route = AsyncMock(return_value=(t.sl_Status.OK,))
 
     packet.source_route = [0x0001, 0x0002]
     await _test_send_packet_unicast(
-        app, packet, options=(t.EmberApsOption.APS_OPTION_RETRY)
+        app,
+        packet,
+        options=(
+            t.EmberApsOption.APS_OPTION_RETRY
+            | t.EmberApsOption.APS_OPTION_ENABLE_ADDRESS_DISCOVERY
+        ),
     )
 
-    assert app._ezsp.setSourceRoute.await_count == 1
-    app._ezsp.setSourceRoute.assert_called_once_with(
-        destination=packet.dst.address,
-        relayList=[0x0001, 0x0002],
+    app._ezsp.set_source_route.assert_called_once_with(
+        nwk=packet.dst.address,
+        relays=[0x0001, 0x0002],
     )
 
 

tests/test_ezsp_v4.py

@@ -337,3 +337,14 @@ async def test_send_broadcast(ezsp_f) -> None:
             messageContents=b"hello",
         )
     ]
+
+
+async def test_source_route(ezsp_f) -> None:
+    ezsp_f.setSourceRoute = AsyncMock(return_value=(t.EmberStatus.SUCCESS,))
+
+    status = await ezsp_f.set_source_route(nwk=0x1234, relays=[0x5678, 0xABCD])
+    assert status == t.sl_Status.OK
+
+    assert ezsp_f.setSourceRoute.mock_calls == [
+        call(destination=0x1234, relayList=[0x5678, 0xABCD])
+    ]

tests/test_ezsp_v9.py

@@ -74,3 +74,13 @@ async def test_write_child_data(ezsp_f) -> None:
             ),
         ),
     ]
+
+
+async def test_source_route(ezsp_f) -> None:
+    ezsp_f.setSourceRoute = AsyncMock(return_value=(t.EmberStatus.SUCCESS,))
+
+    status = await ezsp_f.set_source_route(nwk=0x1234, relays=[0x5678, 0xABCD])
+    assert status == t.sl_Status.OK
+
+    # Nothing happens
+    assert ezsp_f.setSourceRoute.mock_calls == []