Fix unit tests

zigpy · Apr 29, 2024 · e38b6d2 · e38b6d2
1 parent 71ac14c
commit e38b6d2
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 8 deletions.
diff --git a/bellows/ezsp/__init__.py b/bellows/ezsp/__init__.py
@@ -386,7 +386,7 @@ async def _get_nv3_restored_eui64_key(self) -> t.NV3KeyId | None:
  return None
 
  if rsp.status == t.EmberStatus.SUCCESS:
- nv3_restored_eui64, _ = t.EUI64.deserialize(rsp.data)
+ nv3_restored_eui64, _ = t.EUI64.deserialize(rsp.value)
  LOGGER.debug("NV3 restored EUI64: %s=%s", key, nv3_restored_eui64)
 
  return key

diff --git a/bellows/ezsp/protocol.py b/bellows/ezsp/protocol.py
@@ -65,9 +65,9 @@ async def command(self, name, *args) -> Any:
  LOGGER.debug("Send command %s: %s", name, args)
  data = self._ezsp_frame(name, *args)
  self._gw.data(data)
- cmd_id, _, _ = self.COMMANDS[name]
+ cmd_id, _, rx_schema = self.COMMANDS[name]
  future = asyncio.get_running_loop().create_future()
- self._awaiting[self._seq] = (cmd_id, future)
+ self._awaiting[self._seq] = (cmd_id, rx_schema, future)
  self._seq = (self._seq + 1) % 256
 
  async with asyncio_timeout(EZSP_CMD_TIMEOUT):
@@ -100,7 +100,7 @@ def __call__(self, data: bytes) -> None:
  return
 
  try:
- if isinstance(rx_schema, dict):
+ if isinstance(rx_schema, tuple):
  result, data = self.types.deserialize(data, rx_schema)
  else:
  result, data = rx_schema.deserialize(data)
@@ -116,7 +116,7 @@ def __call__(self, data: bytes) -> None:
  LOGGER.debug("Frame contains trailing data: %s", data)
 
  if sequence in self._awaiting:
- expected_id, future = self._awaiting.pop(sequence)
+ expected_id, schema, future = self._awaiting.pop(sequence)
  try:
  if frame_name == "invalidCommand":
  sent_cmd_name = self.COMMANDS_BY_ID[expected_id][0]

diff --git a/bellows/ezsp/v9/commands.py b/bellows/ezsp/v9/commands.py
@@ -1,9 +1,11 @@
+from zigpy.types import Struct, StructField
+
 from . import types as t
 
 
-class GetTokenDataRsp(t.Struct):
+class GetTokenDataRsp(Struct):
  status: t.EmberStatus
- value: t.LVBytes32 = t.StructField(
+ value: t.LVBytes32 = StructField(
  requires=lambda rsp: rsp.status == t.EmberStatus.SUCCESS
  )
 

diff --git a/tests/test_zigbee_repairs.py b/tests/test_zigbee_repairs.py
@@ -7,7 +7,7 @@
 
 from bellows.exception import InvalidCommandError
 from bellows.ezsp import EZSP
-from bellows.ezsp.commands.v9 import GetTokenDataRsp
+from bellows.ezsp.v9.commands import GetTokenDataRsp
 import bellows.types as t
 from bellows.zigbee import repairs