diff --git a/bellows/ezsp/__init__.py b/bellows/ezsp/__init__.py index 2f367812..6892a646 100644 --- a/bellows/ezsp/__init__.py +++ b/bellows/ezsp/__init__.py @@ -532,10 +532,10 @@ def handle_callback(self, *args): except Exception as e: LOGGER.exception("Exception running handler", exc_info=e) - async def set_source_routing(self) -> None: + async def set_source_routing(self, enabled: bool) -> None: """Enable source routing on NCP.""" res = await self.setConcentrator( - on=True, + on=enabled, concentratorType=t.EmberConcentratorType.HIGH_RAM_CONCENTRATOR, minTime=MTOR_MIN_INTERVAL, maxTime=MTOR_MAX_INTERVAL, @@ -548,7 +548,13 @@ async def set_source_routing(self) -> None: LOGGER.warning("Couldn't set concentrator type %s: %s", True, res) if self._ezsp_version >= 8: - await self.setSourceRouteDiscoveryMode(mode=1) + await self.setSourceRouteDiscoveryMode( + mode=( + t.SourceRouteDiscoveryMode.ON + if enabled + else t.SourceRouteDiscoveryMode.OFF + ) + ) def start_ezsp(self): """Mark EZSP as running.""" diff --git a/bellows/types/named.py b/bellows/types/named.py index 83a8cb84..f2a7e897 100644 --- a/bellows/types/named.py +++ b/bellows/types/named.py @@ -2646,3 +2646,11 @@ class SecurityManagerDerivedKeyTypeV13(basic.enum16): # For a TC using hashed link keys, hashed the root key against the supplied EUI in # context. TC_HASHED_LINK_KEY = 5 + + +class SourceRouteDiscoveryMode(basic.enum8): + """Source route discovery mode.""" + + OFF = 0 + ON = 1 + RESCHEDULE = 2 diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py index da0a9551..268bce28 100644 --- a/bellows/zigbee/application.py +++ b/bellows/zigbee/application.py @@ -196,7 +196,9 @@ async def start_network(self): await self._ensure_network_running() if self.config[zigpy.config.CONF_SOURCE_ROUTING]: - await ezsp.set_source_routing() + await ezsp.set_source_routing(enabled=True) + else: + await ezsp.set_source_routing(enabled=False) await ezsp._protocol.update_policies(self.config[CONF_EZSP_POLICIES]) await self.load_network_info(load_devices=False) diff --git a/tests/test_application.py b/tests/test_application.py index bd96a2e7..51d0e299 100644 --- a/tests/test_application.py +++ b/tests/test_application.py @@ -280,6 +280,10 @@ async def mock_leave(*args, **kwargs): t.EmberMulticastTableEntry(multicastId=0x0000, endpoint=0, networkIndex=0), ) proto.setMulticastTableEntry.return_value = [t.EmberStatus.SUCCESS] + proto.setConcentrator.return_value = [t.EmberStatus.SUCCESS] + + if ezsp_version >= 8: + proto.setSourceRouteDiscoveryMode.return_value = [12345] return ezsp_mock diff --git a/tests/test_ezsp.py b/tests/test_ezsp.py index e95abe63..61db3a31 100644 --- a/tests/test_ezsp.py +++ b/tests/test_ezsp.py @@ -407,14 +407,30 @@ async def replacement(command_name, tokenId=None, valueId=None): assert (mfg, brd, ver) == expected -async def test_set_source_routing(ezsp_f): +async def test_set_enable_source_routing(ezsp_f): """Test enabling source routing.""" + ezsp_f.setConcentrator = AsyncMock(return_value=(t.EmberStatus.SUCCESS,)) + ezsp_f.setSourceRouteDiscoveryMode = AsyncMock() + await ezsp_f.set_source_routing(enabled=True) + assert len(ezsp_f.setSourceRouteDiscoveryMode.mock_calls) == 1 + assert ( + ezsp_f.setSourceRouteDiscoveryMode.mock_calls[0].kwargs["mode"] + == t.SourceRouteDiscoveryMode.ON + ) + + +async def test_set_disable_source_routing(ezsp_f): + """Test disabling source routing.""" ezsp_f.setConcentrator = AsyncMock(return_value=(t.EmberStatus.SUCCESS,)) ezsp_f.setSourceRouteDiscoveryMode = AsyncMock() - await ezsp_f.set_source_routing() + await ezsp_f.set_source_routing(enabled=False) assert len(ezsp_f.setSourceRouteDiscoveryMode.mock_calls) == 1 + assert ( + ezsp_f.setSourceRouteDiscoveryMode.mock_calls[0].kwargs["mode"] + == t.SourceRouteDiscoveryMode.OFF + ) async def test_leave_network_error(ezsp_f):