Ensure EZSP repairs work with older firmware

bellows/zigbee/repairs.py

@@ -33,7 +33,7 @@ async def fix_invalid_tclk_partner_ieee(ezsp: EZSP) -> bool:
             t.NV3KeyId.NVM3KEY_STACK_TRUST_CENTER, 0
         )
         assert status == t.EmberStatus.SUCCESS
-    except InvalidCommandError:
+    except (InvalidCommandError, AttributeError):
         LOGGER.warning("NV3 interface not available in this firmware, please upgrade!")
         return False
 

tests/test_ezsp.py

@@ -32,6 +32,23 @@ async def ezsp_f():
         yield api
 
 
+async def make_ezsp(version=4) -> ezsp.EZSP:
+    api = ezsp.EZSP(DEVICE_CONFIG)
+    gw = MagicMock(spec_set=uart.Gateway)
+
+    with patch("bellows.uart.connect", new=AsyncMock(return_value=gw)):
+        await api.connect()
+
+    assert api._ezsp_version == 4
+
+    with patch.object(api, "_command", new=AsyncMock(return_value=[version, 0, 0])):
+        await api.version()
+
+    assert api._ezsp_version == version
+
+    return api
+
+
 async def test_connect(ezsp_f, monkeypatch):
     connected = False
 

tests/test_zigbee_repairs.py

@@ -10,7 +10,7 @@
 import bellows.types as t
 from bellows.zigbee import repairs
 
-from tests.test_ezsp import ezsp_f
+from tests.test_ezsp import ezsp_f, make_ezsp
 
 
 @pytest.fixture
@@ -112,3 +112,62 @@ async def test_fix_invalid_tclk(ezsp_tclk_f: EZSP, caplog) -> None:
             ).serialize(),
         )
     ]
+
+
+@pytest.mark.parametrize("version", EZSP._BY_VERSION)
+async def test_fix_invalid_tclk_all_versions(
+    version: int, ezsp_tclk_f: EZSP, caplog
+) -> None:
+    """Test that the TCLK is fixed (or not) on all versions."""
+
+    ezsp = await make_ezsp(version)
+    fw_has_token_interface = hasattr(ezsp, "setTokenData")
+
+    if fw_has_token_interface:
+        ezsp.setTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
+        ezsp.getTokenData = AsyncMock(
+            return_value=[
+                t.EmberStatus.SUCCESS,
+                t.NV3StackTrustCenterToken(
+                    mode=228,
+                    eui64=t.EmberEUI64.convert("BB:BB:BB:BB:BB:BB:BB:BB"),
+                    key=t.EmberKeyData.convert(
+                        "21:8e:df:b8:50:a0:4a:b6:8b:c6:10:25:bc:4e:93:6a"
+                    ),
+                ).serialize(),
+            ]
+        )
+
+    ezsp.getEui64 = ezsp_tclk_f.getEui64
+    ezsp.getCurrentSecurityState = ezsp_tclk_f.getCurrentSecurityState
+
+    ezsp.getEui64.return_value[0] = t.EmberEUI64.convert("AA:AA:AA:AA:AA:AA:AA:AA")
+    ezsp.getCurrentSecurityState.return_value[
+        1
+    ].trustCenterLongAddress = t.EmberEUI64.convert("BB:BB:BB:BB:BB:BB:BB:BB")
+
+    with caplog.at_level(logging.WARNING):
+        assert (
+            await repairs.fix_invalid_tclk_partner_ieee(ezsp) is fw_has_token_interface
+        )
+
+    assert "Fixing invalid TCLK" in caplog.text
+
+    if fw_has_token_interface:
+        assert "NV3 interface not available in this firmware" not in caplog.text
+
+        assert ezsp.setTokenData.mock_calls == [
+            call(
+                t.NV3KeyId.NVM3KEY_STACK_TRUST_CENTER,
+                0,
+                t.NV3StackTrustCenterToken(
+                    mode=228,
+                    eui64=t.EmberEUI64.convert("AA:AA:AA:AA:AA:AA:AA:AA"),
+                    key=t.EmberKeyData.convert(
+                        "21:8e:df:b8:50:a0:4a:b6:8b:c6:10:25:bc:4e:93:6a"
+                    ),
+                ).serialize(),
+            )
+        ]
+    else:
+        assert "NV3 interface not available in this firmware" in caplog.text