Fix existing unit tests

bellows/ezsp/__init__.py

@@ -156,16 +156,17 @@ async def reset(self):
 
     def _switch_protocol_version(self, version: int) -> None:
         LOGGER.debug("Switching to EZSP protocol version %d", version)
+        self._ezsp_version = version
 
         if version not in self._BY_VERSION:
             LOGGER.warning(
                 "Protocol version %s is not supported, using version %s instead",
                 version,
                 EZSP_LATEST,
             )
+            # We replace the protocol object but keep the version correct
             version = EZSP_LATEST
 
-        self._ezsp_version = version
         self._protocol = self._BY_VERSION[version](self.handle_callback, self._gw)
 
     async def version(self):

bellows/ezsp/protocol.py

@@ -55,7 +55,9 @@ def _ezsp_frame_tx(self, name: str) -> bytes:
     async def pre_permit(self, time_s: int) -> None:
         """Schedule task before allowing new joins."""
 
-    async def add_transient_link_key(self, ieee: t.EUI64, key: t.KeyData) -> None:
+    async def add_transient_link_key(
+        self, ieee: t.EUI64, key: t.KeyData
+    ) -> t.EmberStatus:
         """Add a transient link key."""
 
     async def command(self, name, *args) -> Any:

bellows/ezsp/v13/__init__.py

@@ -23,9 +23,11 @@ class EZSPv13(EZSPv12):
 
     async def add_transient_link_key(
         self, ieee: t.EUI64, key: t.KeyData
-    ) -> tuple[t.EmberStatus]:
-        return await self.importTransientKey(
+    ) -> t.EmberStatus:
+        (status,) = await self.importTransientKey(
             ieee,
             key,
             v13_types.sl_zb_sec_man_flags_t.NONE,
         )
+
+        return status

bellows/ezsp/v5/__init__.py

@@ -37,8 +37,9 @@ def _ezsp_frame_rx(self, data: bytes) -> tuple[int, int, bytes]:
 
     async def add_transient_link_key(
         self, ieee: t.EUI64, key: t.KeyData
-    ) -> tuple[t.EmberStatus]:
-        return await self.addTransientLinkKey(ieee, key)
+    ) -> t.EmberStatus:
+        (status,) = await self.addTransientLinkKey(ieee, key)
+        return status
 
     async def pre_permit(self, time_s: int) -> None:
         """Add pre-shared TC Link key."""

bellows/zigbee/application.py

@@ -924,10 +924,10 @@ async def permit_with_link_key(
     ) -> None:
         """Permits a new device to join with the given IEEE and link key."""
 
-        v = await self._ezsp.add_transient_link_key(node, link_key)
+        status = await self._ezsp.add_transient_link_key(node, link_key)
 
-        if v[0] != t.EmberStatus.SUCCESS:
-            raise Exception("Failed to set link key")
+        if status != t.EmberStatus.SUCCESS:
+            raise ControllerError("Failed to set link key")
 
         if self._ezsp.ezsp_version >= 8:
             await self._ezsp.setPolicy(

tests/test_application.py

@@ -71,6 +71,7 @@ def ezsp_mock(ieee):
     mock_ezsp.wait_for_stack_status.return_value.__enter__ = AsyncMock(
         return_value=t.EmberStatus.NETWORK_UP
     )
+    mock_ezsp.add_transient_link_key = AsyncMock(return_value=t.EmberStatus.SUCCESS)
     mock_ezsp._protocol = AsyncMock()
 
     type(mock_ezsp).types = ezsp_t7
@@ -615,49 +616,33 @@ async def test_permit_ncp(app):
     assert app._ezsp.permitJoining.call_count == 1
 
 
-@pytest.mark.parametrize(
-    "version, tc_policy_count, ezsp_types",
-    ((4, 0, t), (5, 0, ezsp_t5), (6, 0, ezsp_t6), (7, 0, ezsp_t7), (8, 1, ezsp_t8)),
-)
-async def test_permit_with_link_key_ieee(
-    app, ieee, version, tc_policy_count, ezsp_types
-):
-    p1 = patch("zigpy.application.ControllerApplication.permit")
-    p2 = patch.object(app._ezsp, "types", ezsp_types)
-
-    with patch.object(app._ezsp, "ezsp_version", version), p1 as permit_mock, p2:
+async def test_permit_with_link_key(app, ieee):
+    with patch("zigpy.application.ControllerApplication.permit") as permit_mock:
         await app.permit_with_link_key(
             ieee,
-            zigpy_t.KeyData.convert("11:22:33:44:55:66:77:88:11:22:33:44:55:66:77:88:"),
+            zigpy_t.KeyData.convert("11:22:33:44:55:66:77:88:11:22:33:44:55:66:77:88"),
             60,
         )
 
-    assert app._ezsp.addTransientLinkKey.await_count == 1
     assert permit_mock.await_count == 1
-    assert app._ezsp.setPolicy.await_count == tc_policy_count
+    assert app._ezsp.add_transient_link_key.await_count == 1
 
 
-async def test_permit_with_link_key_failed_add_key(app, ieee):
-    app._ezsp.addTransientLinkKey = AsyncMock(return_value=[1, 1])
+async def test_permit_with_link_key_failure(app, ieee):
+    app._ezsp.add_transient_link_key.return_value = t.EmberStatus.ERR_FATAL
 
-    with pytest.raises(Exception):
-        await app.permit_with_link_key(
-            ieee,
-            zigpy_t.KeyData.convert("11:22:33:44:55:66:77:88:11:22:33:44:55:66:77:88:"),
-            60,
-        )
+    with patch("zigpy.application.ControllerApplication.permit") as permit_mock:
+        with pytest.raises(ControllerError):
+            await app.permit_with_link_key(
+                ieee,
+                zigpy_t.KeyData.convert(
+                    "11:22:33:44:55:66:77:88:11:22:33:44:55:66:77:88"
+                ),
+                60,
+            )
 
-
-async def test_permit_with_link_key_failed_set_policy(app, ieee):
-    app._ezsp.addTransientLinkKey = AsyncMock(return_value=[0])
-    app._ezsp.setPolicy = AsyncMock(return_value=[1])
-
-    with pytest.raises(Exception):
-        await app.permit_with_link_key(
-            ieee,
-            zigpy_t.KeyData.convert("11:22:33:44:55:66:77:88:11:22:33:44:55:66:77:88:"),
-            60,
-        )
+    assert permit_mock.await_count == 0
+    assert app._ezsp.add_transient_link_key.await_count == 1
 
 
 @pytest.fixture
@@ -775,33 +760,6 @@ async def test_send_packet_unicast_source_route_ezsp7(make_app, packet):
     )
 
 
-async def test_send_packet_unicast_source_route_ezsp8_have_relays(make_app, packet):
-    app = make_app({zigpy.config.CONF_SOURCE_ROUTING: True})
-    app._ezsp.ezsp_version = 8
-
-    device = MagicMock()
-    device.relays = [0x0003]
-
-    app.get_device = MagicMock(return_value=device)
-
-    packet.source_route = [0x0001, 0x0002]
-    await _test_send_packet_unicast(app, packet)
-
-    aps_frame = app._ezsp.sendUnicast.mock_calls[0].args[2]
-    assert t.EmberApsOption.APS_OPTION_ENABLE_ROUTE_DISCOVERY not in aps_frame.options
-
-
-async def test_send_packet_unicast_source_route_ezsp8_no_relays(make_app, packet):
-    app = make_app({zigpy.config.CONF_SOURCE_ROUTING: True})
-    app._ezsp.ezsp_version = 8
-
-    packet.source_route = [0x0001, 0x0002]
-    await _test_send_packet_unicast(app, packet)
-
-    aps_frame = app._ezsp.sendUnicast.mock_calls[0].args[2]
-    assert t.EmberApsOption.APS_OPTION_ENABLE_ROUTE_DISCOVERY in aps_frame.options
-
-
 async def test_send_packet_unicast_retries_success(app, packet):
     await _test_send_packet_unicast(
         app,
@@ -1585,32 +1543,6 @@ async def test_startup_new_coordinator_no_groups_joined(app, ieee):
     p2.assert_not_called()
 
 
-@pytest.mark.parametrize("enable_source_routing", [True, False])
-async def test_startup_source_routing(make_app, ieee, enable_source_routing):
-    """Existing relays are cleared on startup."""
-
-    app = make_app({zigpy.config.CONF_SOURCE_ROUTING: enable_source_routing})
-
-    app._ezsp.ezsp_version = 9
-    app._ezsp.update_policies = AsyncMock()
-
-    app._ensure_network_running = AsyncMock()
-    app.load_network_info = AsyncMock()
-    app.state.node_info.ieee = ieee
-
-    mock_device = MagicMock()
-    mock_device.relays = sentinel.relays
-    mock_device.initialize = AsyncMock()
-    app.devices[0xABCD] = mock_device
-
-    await app.start_network()
-
-    if enable_source_routing:
-        assert mock_device.relays is None
-    else:
-        assert mock_device.relays is sentinel.relays
-
-
 @pytest.mark.parametrize(
     "scan_results",
     [

tests/test_application_network_state.py

@@ -102,7 +102,17 @@ def _mock_app_for_load(app):
 
     ezsp.getNodeId = AsyncMock(return_value=[t.EmberNodeId(0x0000)])
     ezsp.getEui64 = AsyncMock(return_value=[t.EUI64.convert("00:12:4b:00:1c:a1:b8:46")])
-    ezsp.getConfigurationValue = AsyncMock(return_value=[t.EmberStatus.SUCCESS, 5])
+
+    def get_configuration_value(config_id):
+        size = {
+            app._ezsp.types.EzspConfigId.CONFIG_ADDRESS_TABLE_SIZE: t.uint8_t(20),
+            app._ezsp.types.EzspConfigId.CONFIG_KEY_TABLE_SIZE: t.uint8_t(13),
+            app._ezsp.types.EzspConfigId.CONFIG_SECURITY_LEVEL: t.uint8_t(5),
+        }[config_id]
+
+        return [app._ezsp.types.EmberStatus.SUCCESS, size]
+
+    ezsp.getConfigurationValue = AsyncMock(side_effect=get_configuration_value)
 
     def get_key(key_type):
         key = {
@@ -330,10 +340,25 @@ def get_addr_table_eui64(index):
 
 def _mock_app_for_write(app, network_info, node_info, ezsp_ver=None):
     ezsp = app._ezsp
-    ezsp.networkState = AsyncMock(
-        return_value=[ezsp.types.EmberNetworkStatus.JOINED_NETWORK]
-    )
-    ezsp.leaveNetwork = AsyncMock(return_value=[t.EmberStatus.NETWORK_DOWN])
+
+    network_state = ezsp.types.EmberNetworkStatus.JOINED_NETWORK
+    ezsp.networkState = AsyncMock(side_effect=lambda: [network_state])
+
+    def leave_network():
+        nonlocal network_state
+        network_state = ezsp.types.EmberNetworkStatus.NO_NETWORK
+
+        return [t.EmberStatus.NETWORK_DOWN]
+
+    def form_network(params):
+        nonlocal network_state
+        network_state = ezsp.types.EmberNetworkStatus.JOINED_NETWORK
+
+        return [t.EmberStatus.SUCCESS]
+
+    ezsp.leaveNetwork = AsyncMock(side_effect=leave_network)
+    ezsp.formNetwork = AsyncMock(side_effect=form_network)
+
     ezsp.getEui64 = AsyncMock(return_value=[t.EUI64.convert("00:12:4b:00:1c:a1:b8:46")])
 
     ezsp.setInitialSecurityState = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
@@ -361,7 +386,6 @@ def _mock_app_for_write(app, network_info, node_info, ezsp_ver=None):
         else:
             ezsp.setValue = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
 
-    ezsp.formNetwork = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
     ezsp.setValue = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
     ezsp.setMfgToken = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
     ezsp.getTokenData = AsyncMock(return_value=[t.EmberStatus.LIBRARY_NOT_PRESENT, b""])

tests/test_ezsp_v10.py

@@ -28,7 +28,11 @@ def test_ezsp_frame_rx(ezsp_f):
 async def test_pre_permit(ezsp_f):
     """Test pre permit."""
     p1 = patch.object(ezsp_f, "setPolicy", new=AsyncMock())
-    p2 = patch.object(ezsp_f, "addTransientLinkKey", new=AsyncMock())
+    p2 = patch.object(
+        ezsp_f,
+        "addTransientLinkKey",
+        new=AsyncMock(return_value=[ezsp_f.types.EmberStatus.SUCCESS]),
+    )
     with p1 as pre_permit_mock, p2 as tclk_mock:
         await ezsp_f.pre_permit(-1.9)
     assert pre_permit_mock.await_count == 2
@@ -266,6 +270,7 @@ def test_command_frames(ezsp_f):
     "setRoutingShortcutThreshold": 0x00D0,
     "setSecurityKey": 0x00CA,
     "setSecurityParameters": 0x00CB,
+    "setSourceRoute": 0x00AE,
     "setSourceRouteDiscoveryMode": 0x005A,
     "setTimer": 0x000E,
     "setToken": 0x0009,

tests/test_ezsp_v11.py

@@ -28,7 +28,11 @@ def test_ezsp_frame_rx(ezsp_f):
 async def test_pre_permit(ezsp_f):
     """Test pre permit."""
     p1 = patch.object(ezsp_f, "setPolicy", new=AsyncMock())
-    p2 = patch.object(ezsp_f, "addTransientLinkKey", new=AsyncMock())
+    p2 = patch.object(
+        ezsp_f,
+        "addTransientLinkKey",
+        new=AsyncMock(return_value=[ezsp_f.types.EmberStatus.SUCCESS]),
+    )
     with p1 as pre_permit_mock, p2 as tclk_mock:
         await ezsp_f.pre_permit(-1.9)
     assert pre_permit_mock.await_count == 2
@@ -266,6 +270,7 @@ def test_command_frames(ezsp_f):
     "setRoutingShortcutThreshold": 0x00D0,
     "setSecurityKey": 0x00CA,
     "setSecurityParameters": 0x00CB,
+    "setSourceRoute": 0x00AE,
     "setSourceRouteDiscoveryMode": 0x005A,
     "setTimer": 0x000E,
     "setToken": 0x0009,

tests/test_ezsp_v12.py

@@ -28,7 +28,11 @@ def test_ezsp_frame_rx(ezsp_f):
 async def test_pre_permit(ezsp_f):
     """Test pre permit."""
     p1 = patch.object(ezsp_f, "setPolicy", new=AsyncMock())
-    p2 = patch.object(ezsp_f, "addTransientLinkKey", new=AsyncMock())
+    p2 = patch.object(
+        ezsp_f,
+        "addTransientLinkKey",
+        new=AsyncMock(return_value=[ezsp_f.types.EmberStatus.SUCCESS]),
+    )
     with p1 as pre_permit_mock, p2 as tclk_mock:
         await ezsp_f.pre_permit(-1.9)
     assert pre_permit_mock.await_count == 2
@@ -280,6 +284,7 @@ def test_command_frames(ezsp_f):
     "setRoutingShortcutThreshold": 0x00D0,
     "setSecurityKey": 0x00CA,
     "setSecurityParameters": 0x00CB,
+    "setSourceRoute": 0x00AE,
     "setSourceRouteDiscoveryMode": 0x005A,
     "setTimer": 0x000E,
     "setToken": 0x0009,

tests/test_ezsp_v5.py

@@ -27,7 +27,11 @@ def test_ezsp_frame_rx(ezsp_f):
 
 async def test_pre_permit(ezsp_f):
     """Test pre permit."""
-    p2 = patch.object(ezsp_f, "addTransientLinkKey", new=AsyncMock())
+    p2 = patch.object(
+        ezsp_f,
+        "addTransientLinkKey",
+        new=AsyncMock(return_value=[ezsp_f.types.EmberStatus.SUCCESS]),
+    )
     with p2 as tclk_mock:
         await ezsp_f.pre_permit(1)
     assert tclk_mock.await_count == 1

tests/test_ezsp_v8.py

@@ -28,7 +28,11 @@ def test_ezsp_frame_rx(ezsp_f):
 async def test_pre_permit(ezsp_f):
     """Test pre permit."""
     p1 = patch.object(ezsp_f, "setPolicy", new=AsyncMock())
-    p2 = patch.object(ezsp_f, "addTransientLinkKey", new=AsyncMock())
+    p2 = patch.object(
+        ezsp_f,
+        "addTransientLinkKey",
+        new=AsyncMock(return_value=[ezsp_f.types.EmberStatus.SUCCESS]),
+    )
     with p1 as pre_permit_mock, p2 as tclk_mock:
         await ezsp_f.pre_permit(-1.9)
     assert pre_permit_mock.await_count == 2
@@ -258,6 +262,7 @@ def test_command_frames(ezsp_f):
     "setRoutingShortcutThreshold": 0x00D0,
     "setSecurityKey": 0x00CA,
     "setSecurityParameters": 0x00CB,
+    "setSourceRoute": 0x00AE,
     "setSourceRouteDiscoveryMode": 0x005A,
     "setTimer": 0x000E,
     "setToken": 0x0009,

tests/test_ezsp_v9.py

@@ -28,7 +28,11 @@ def test_ezsp_frame_rx(ezsp_f):
 async def test_pre_permit(ezsp_f):
     """Test pre permit."""
     p1 = patch.object(ezsp_f, "setPolicy", new=AsyncMock())
-    p2 = patch.object(ezsp_f, "addTransientLinkKey", new=AsyncMock())
+    p2 = patch.object(
+        ezsp_f,
+        "addTransientLinkKey",
+        new=AsyncMock(return_value=[ezsp_f.types.EmberStatus.SUCCESS]),
+    )
     with p1 as pre_permit_mock, p2 as tclk_mock:
         await ezsp_f.pre_permit(-1.9)
     assert pre_permit_mock.await_count == 2
@@ -266,6 +270,7 @@ def test_command_frames(ezsp_f):
     "setRoutingShortcutThreshold": 0x00D0,
     "setSecurityKey": 0x00CA,
     "setSecurityParameters": 0x00CB,
+    "setSourceRoute": 0x00AE,
     "setSourceRouteDiscoveryMode": 0x005A,
     "setTimer": 0x000E,
     "setToken": 0x0009,