Fix uses of getTokenData

bellows/ezsp/__init__.py

@@ -379,14 +379,14 @@ async def _get_nv3_restored_eui64_key(self) -> t.NV3KeyId | None:
  t.NV3KeyId.NVM3KEY_STACK_RESTORED_EUI64, # RCP firmware
  ):
  try:
- status, data = await self.getTokenData(key, 0)
+ rsp = await self.getTokenData(key, 0)
  except (InvalidCommandError, AttributeError):
  # Either the command doesn't exist in the EZSP version, or the command
  # is not implemented in the firmware
  return None
 
- if status == t.EmberStatus.SUCCESS:
- nv3_restored_eui64, _ = t.EUI64.deserialize(data)
+ if rsp.status == t.EmberStatus.SUCCESS:
+ nv3_restored_eui64, _ = t.EUI64.deserialize(rsp.data)
  LOGGER.debug("NV3 restored EUI64: %s=%s", key, nv3_restored_eui64)
 
  return key

bellows/ezsp/v9/commands.py

@@ -3,7 +3,7 @@
 
 class GetTokenDataRsp(t.Struct):
  status: t.EmberStatus
- data: t.LVBytes32 = t.StructField(
+ value: t.LVBytes32 = t.StructField(
  requires=lambda rsp: rsp.status == t.EmberStatus.SUCCESS
  )
 

bellows/zigbee/repairs.py

@@ -29,15 +29,13 @@ async def fix_invalid_tclk_partner_ieee(ezsp: EZSP) -> bool:
  )
 
  try:
- (status, value) = await ezsp.getTokenData(
- t.NV3KeyId.NVM3KEY_STACK_TRUST_CENTER, 0
- )
- assert status == t.EmberStatus.SUCCESS
+ rsp = await ezsp.getTokenData(t.NV3KeyId.NVM3KEY_STACK_TRUST_CENTER, 0)
+ assert rsp.status == t.EmberStatus.SUCCESS
  except (InvalidCommandError, AttributeError, AssertionError):
  LOGGER.warning("NV3 interface not available in this firmware, please upgrade!")
  return False
 
- token, remaining = t.NV3StackTrustCenterToken.deserialize(value)
+ token, remaining = t.NV3StackTrustCenterToken.deserialize(rsp.value)
  assert not remaining
  assert token.eui64 == state.trustCenterLongAddress
 

tests/test_application.py

@@ -18,6 +18,7 @@
 import bellows.ezsp.v6.types as ezsp_t6
 import bellows.ezsp.v7.types as ezsp_t7
 import bellows.ezsp.v8.types as ezsp_t8
+from bellows.ezsp.v9.commands import GetTokenDataRsp
 import bellows.types.struct
 import bellows.uart as uart
 import bellows.zigbee.application
@@ -154,7 +155,9 @@ async def nop_mock():
  ezsp_mock.readAndClearCounters = AsyncMock(side_effect=nop_mock)
  ezsp_mock._protocol = AsyncMock()
  ezsp_mock.setConcentrator = AsyncMock()
- ezsp_mock.getTokenData = AsyncMock(return_value=[t.EmberStatus.ERR_FATAL, b""])
+ ezsp_mock.getTokenData = AsyncMock(
+ return_value=GetTokenDataRsp(status=t.EmberStatus.ERR_FATAL)
+ )
  ezsp_mock._command = AsyncMock(return_value=t.EmberStatus.SUCCESS)
  ezsp_mock.addEndpoint = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
  ezsp_mock.setConfigurationValue = AsyncMock(return_value=[t.EmberStatus.SUCCESS])

tests/test_application_network_state.py

@@ -8,6 +8,7 @@
 
 from bellows.exception import EzspError
 from bellows.ezsp import EZSP
+from bellows.ezsp.v9.commands import GetTokenDataRsp
 import bellows.types as t
 
 from tests.async_mock import AsyncMock, PropertyMock
@@ -518,7 +519,9 @@ def form_network(params):
 
  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""])
+ ezsp.getTokenData = AsyncMock(
+ return_value=GetTokenDataRsp(status=t.EmberStatus.LIBRARY_NOT_PRESENT)
+ )
 
 
 @pytest.mark.parametrize("ezsp_ver", [4, 7, 13])

tests/test_ezsp.py

@@ -18,6 +18,8 @@
 
 from unittest.mock import ANY, AsyncMock, MagicMock, call, patch, sentinel
 
+from bellows.ezsp.v9.commands import GetTokenDataRsp
+
 DEVICE_CONFIG = {
  zigpy.config.CONF_DEVICE_PATH: "/dev/null",
  zigpy.config.CONF_DEVICE_BAUDRATE: 115200,
@@ -524,9 +526,9 @@ async def test_can_rewrite_custom_eui64(ezsp_f, tokens, expected_key, expected_r
 
  def get_token_data(key, index):
  if key not in tokens or index != 0:
- return [t.EmberStatus.ERR_FATAL, b""]
+ return GetTokenDataRsp(status=t.EmberStatus.ERR_FATAL)
 
- return [t.EmberStatus.SUCCESS, tokens[key]]
+ return GetTokenDataRsp(status=t.EmberStatus.SUCCESS, value=tokens[key])
 
  ezsp_f.getTokenData = AsyncMock(side_effect=get_token_data)
 
@@ -623,7 +625,9 @@ async def test_write_custom_eui64_rcp(ezsp_f):
 
  # RCP firmware does not support manufacturing tokens
  ezsp_f.getMfgToken = AsyncMock(return_value=[b""])
- ezsp_f.getTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS, b"\xFF" * 8])
+ ezsp_f.getTokenData = AsyncMock(
+ return_value=GetTokenDataRsp(status=t.EmberStatus.SUCCESS, value=b"\xFF" * 8)
+ )
 
  await ezsp_f.write_custom_eui64(new_eui64)
 
@@ -858,7 +862,9 @@ async def test_reset_custom_eui64(ezsp_f):
  assert len(ezsp_f.setTokenData.mock_calls) == 0
 
  # With NV3 interface
- ezsp_f.getTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS, b"\xAB" * 8])
+ ezsp_f.getTokenData = AsyncMock(
+ return_value=GetTokenDataRsp(status=t.EmberStatus.SUCCESS, value=b"\xAB" * 8)
+ )
  await ezsp_f.reset_custom_eui64()
  assert ezsp_f.setTokenData.mock_calls == [
  call(t.NV3KeyId.CREATOR_STACK_RESTORED_EUI64, 0, t.LVBytes32(b"\xFF" * 8))

tests/test_zigbee_repairs.py

@@ -7,6 +7,7 @@
 
 from bellows.exception import InvalidCommandError
 from bellows.ezsp import EZSP
+from bellows.ezsp.commands.v9 import GetTokenDataRsp
 import bellows.types as t
 from bellows.zigbee import repairs
 
@@ -71,16 +72,16 @@ async def test_fix_invalid_tclk(ezsp_tclk_f: EZSP, caplog) -> None:
 
  ezsp_tclk_f.setTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
  ezsp_tclk_f.getTokenData = AsyncMock(
- return_value=[
- t.EmberStatus.SUCCESS,
- t.NV3StackTrustCenterToken(
+ return_value=GetTokenDataRsp(
+ status=t.EmberStatus.SUCCESS,
+ value=t.NV3StackTrustCenterToken(
  mode=228,
  eui64=t.EUI64.convert("BB:BB:BB:BB:BB:BB:BB:BB"),
  key=t.KeyData.convert(
  "21:8e:df:b8:50:a0:4a:b6:8b:c6:10:25:bc:4e:93:6a"
  ),
  ).serialize(),
- ]
+ )
  )
  ezsp_tclk_f.getEui64.return_value[0] = t.EUI64.convert("AA:AA:AA:AA:AA:AA:AA:AA")
  ezsp_tclk_f.getCurrentSecurityState.return_value[
@@ -121,21 +122,23 @@ async def test_fix_invalid_tclk_all_versions(
  if fw_has_token_interface:
  ezsp.setTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
  ezsp.getTokenData = AsyncMock(
- return_value=[
- t.EmberStatus.SUCCESS,
- t.NV3StackTrustCenterToken(
+ return_value=GetTokenDataRsp(
+ status=t.EmberStatus.SUCCESS,
+ value=t.NV3StackTrustCenterToken(
  mode=228,
  eui64=t.EUI64.convert("BB:BB:BB:BB:BB:BB:BB:BB"),
  key=t.KeyData.convert(
  "21:8e:df:b8:50:a0:4a:b6:8b:c6:10:25:bc:4e:93:6a"
  ),
  ).serialize(),
- ]
+ )
  )
 
  if not has_library:
  ezsp.setTokenData = AsyncMock(return_value=[t.EmberStatus.LIBRARY_NOT_LOADED])
- ezsp.getTokenData = AsyncMock(return_value=[t.EmberStatus.LIBRARY_NOT_LOADED])
+ ezsp.getTokenData = AsyncMock(
+ return_value=GetTokenDataRsp(status=t.EmberStatus.LIBRARY_NOT_LOADED)
+ )
 
  ezsp.getEui64 = ezsp_tclk_f.getEui64
  ezsp.getCurrentSecurityState = ezsp_tclk_f.getCurrentSecurityState