Repository: zigpy/bellows

Issue:
Error doing job: Exception in callback ThreadsafeProxy
(Hoping I'm creating the issue in the right place, if not, please direct me to the correct repo.)
Since testing out home assistant beta 2024.5, I'm getting the following error:


```
Logger: homeassistant
Source: /usr/src/homeassistant/homeassistant/runner.py:147
First occurred: 08:55:34 (39 occurrences)
Last logged: 09:14:19

Error doing job: Exception in callback ThreadsafeProxy.__getattr__.<locals>.func_wrapper.<locals>.check_result_wrapper() at /usr/local/lib/python3.12/site-packages/bellows/thread.py:110
Traceback (most recent call last):
  File "/usr/local/lib/python3.12/asyncio/events.py", line 88, in _run
    self._context.run(self._callback, *self._args)
  File "/usr/local/lib/python3.12/site-packages/bellows/thread.py", line 111, in check_result_wrapper
    result = call()
             ^^^^^^
  File "/usr/local/lib/python3.12/site-packages/bellows/ezsp/__init__.py", line 330, in frame_received
    self._protocol(data)
  File "/usr/local/lib/python3.12/site-packages/bellows/ezsp/protocol.py", line 103, in __call__
    result, data = self.types.deserialize(data, schema)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/bellows/types/__init__.py", line 9, in deserialize
    value, data = type_.deserialize(data)
                  ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/zigpy/types/basic.py", line 197, in deserialize
    raise ValueError(f"Data is too short to contain {byte_size} bytes")
ValueError: Data is too short to contain 2 bytes
```


This error is accompanied by the following warning:

```
Logger: bellows.ezsp.protocol
Source: runner.py:190
First occurred: 08:55:34 (39 occurrences)
Last logged: 09:14:19

Failed to parse frame changeSourceRouteHandler: b'0c6f15'
Failed to parse frame changeSourceRouteHandler: b'0c4794'
Failed to parse frame changeSourceRouteHandler: b'0ccc82'
```

turned on debug mode and captured the following log from startup: https://dpaste.org/h98Kv

Discussion:
@puddly: Could you enable debug mode for the ZHA integration and reload it? @Yardco: Sure thing (added the log directly here, struggeling with adding it to dpaste)
[log file](https://github.com/zigpy/bellows/files/15149925/home-assistant_zha_2024-04-29T13-32-55.233Z.log)


Solution patch:
diff --git a/bellows/ezsp/__init__.py b/bellows/ezsp/__init__.py
index 1872a74..4aee11c 100644
--- a/bellows/ezsp/__init__.py
+++ b/bellows/ezsp/__init__.py
@@ -379,14 +379,14 @@ class EZSP:
             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.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
index 9518ba8..ef7157a 100644
--- a/bellows/ezsp/protocol.py
+++ b/bellows/ezsp/protocol.py
@@ -65,9 +65,9 @@ class ProtocolHandler(abc.ABC):
         LOGGER.debug("Send command %s: %s", name, args)
         data = self._ezsp_frame(name, *args)
         self._gw.data(data)
-        c = self.COMMANDS[name]
-        future = asyncio.Future()
-        self._awaiting[self._seq] = (c[0], c[2], future)
+        cmd_id, _, rx_schema = self.COMMANDS[name]
+        future = asyncio.get_running_loop().create_future()
+        self._awaiting[self._seq] = (cmd_id, rx_schema, future)
         self._seq = (self._seq + 1) % 256
 
         async with asyncio_timeout(EZSP_CMD_TIMEOUT):
@@ -89,7 +89,7 @@ class ProtocolHandler(abc.ABC):
         sequence, frame_id, data = self._ezsp_frame_rx(data)
 
         try:
-            frame_name, _, schema = self.COMMANDS_BY_ID[frame_id]
+            frame_name, _, rx_schema = self.COMMANDS_BY_ID[frame_id]
         except KeyError:
             LOGGER.warning(
                 "Unknown application frame 0x%04X received: %s (%s).  This is a bug!",
@@ -100,7 +100,10 @@ class ProtocolHandler(abc.ABC):
             return
 
         try:
-            result, data = self.types.deserialize(data, schema)
+            if isinstance(rx_schema, tuple):
+                result, data = self.types.deserialize(data, rx_schema)
+            else:
+                result, data = rx_schema.deserialize(data)
         except Exception:
             LOGGER.warning(
                 "Failed to parse frame %s: %s", frame_name, binascii.hexlify(data)
diff --git a/bellows/ezsp/v10/commands.py b/bellows/ezsp/v10/commands.py
index b70c529..34c555c 100644
--- a/bellows/ezsp/v10/commands.py
+++ b/bellows/ezsp/v10/commands.py
@@ -1,3 +1,5 @@
+from bellows.ezsp.v9.commands import GetTokenDataRsp
+
 from . import types as t
 
 COMMANDS = {
@@ -691,7 +693,7 @@ COMMANDS = {
     "getTokenData": (
         0x0102,
         (t.uint32_t, t.uint32_t),
-        (t.EmberStatus, t.LVBytes32),
+        GetTokenDataRsp,
     ),
     "setTokenData": (
         0x0103,
diff --git a/bellows/ezsp/v9/commands.py b/bellows/ezsp/v9/commands.py
index ee92562..21188ca 100644
--- a/bellows/ezsp/v9/commands.py
+++ b/bellows/ezsp/v9/commands.py
@@ -1,5 +1,15 @@
+from zigpy.types import Struct, StructField
+
 from . import types as t
 
+
+class GetTokenDataRsp(Struct):
+    status: t.EmberStatus
+    value: t.LVBytes32 = StructField(
+        requires=lambda rsp: rsp.status == t.EmberStatus.SUCCESS
+    )
+
+
 COMMANDS = {
     # 4. Configuration frames
     "version": (0x0000, (t.uint8_t,), (t.uint8_t, t.uint8_t, t.uint16_t)),
@@ -687,7 +697,7 @@ COMMANDS = {
     "getTokenData": (
         0x0102,
         (t.uint32_t, t.uint32_t),
-        (t.EmberStatus, t.LVBytes32),
+        GetTokenDataRsp,
     ),
     "setTokenData": (
         0x0103,
diff --git a/bellows/zigbee/repairs.py b/bellows/zigbee/repairs.py
index ab01f15..671f864 100644
--- a/bellows/zigbee/repairs.py
+++ b/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
 


Test patch:
diff --git a/tests/test_application.py b/tests/test_application.py
index 7434e23..ebcb32d 100644
--- a/tests/test_application.py
+++ b/tests/test_application.py
@@ -18,6 +18,7 @@ import bellows.ezsp.v5.types as ezsp_t5
 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 @@ def _create_app_for_startup(
     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])
diff --git a/tests/test_application_network_state.py b/tests/test_application_network_state.py
index 0b02074..0689e78 100644
--- a/tests/test_application_network_state.py
+++ b/tests/test_application_network_state.py
@@ -8,6 +8,7 @@ import zigpy.zdo.types as zdo_t
 
 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 _mock_app_for_write(app, network_info, node_info, ezsp_ver=7):
 
     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])
diff --git a/tests/test_ezsp.py b/tests/test_ezsp.py
index 6121d84..e76b118 100644
--- a/tests/test_ezsp.py
+++ b/tests/test_ezsp.py
@@ -18,6 +18,8 @@ else:
 
 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))
diff --git a/tests/test_ezsp_protocol.py b/tests/test_ezsp_protocol.py
index 6b72775..2979451 100644
--- a/tests/test_ezsp_protocol.py
+++ b/tests/test_ezsp_protocol.py
@@ -6,6 +6,9 @@ import pytest
 from bellows.ezsp import EZSP
 import bellows.ezsp.v4
 import bellows.ezsp.v4.types as t
+import bellows.ezsp.v9
+from bellows.ezsp.v9.commands import GetTokenDataRsp
+from bellows.types import NV3KeyId
 
 from .async_mock import ANY, AsyncMock, MagicMock, call, patch
 
@@ -16,6 +19,12 @@ def prot_hndl():
     return bellows.ezsp.v4.EZSPv4(MagicMock(), MagicMock())
 
 
+@pytest.fixture
+def prot_hndl_v9():
+    """Protocol handler mock."""
+    return bellows.ezsp.v9.EZSPv9(MagicMock(), MagicMock())
+
+
 async def test_command(prot_hndl):
     coro = prot_hndl.command("nop")
     asyncio.get_running_loop().call_soon(
@@ -94,3 +103,21 @@ async def test_logging_frame_parsing_failure(prot_hndl, caplog) -> None:
             prot_hndl(b"\xAA\xAA\x71\x22")
 
         assert "Failed to parse frame getKeyTableEntry: b'22'" in caplog.text
+
+
+async def test_parsing_schema_response(prot_hndl_v9):
+    """Test parsing data with a struct schema."""
+
+    coro = prot_hndl_v9.command(
+        "getTokenData", NV3KeyId.CREATOR_STACK_RESTORED_EUI64, 0
+    )
+    asyncio.get_running_loop().call_soon(
+        lambda: prot_hndl_v9(
+            bytes([prot_hndl_v9._seq - 1, 0x00, 0x00])
+            + t.uint16_t(prot_hndl_v9.COMMANDS["getTokenData"][0]).serialize()
+            + bytes([0xB5])
+        )
+    )
+
+    rsp = await coro
+    assert rsp == GetTokenDataRsp(status=t.EmberStatus.LIBRARY_NOT_PRESENT)
diff --git a/tests/test_zigbee_repairs.py b/tests/test_zigbee_repairs.py
index 8ecc3b4..6a4a03f 100644
--- a/tests/test_zigbee_repairs.py
+++ b/tests/test_zigbee_repairs.py
@@ -7,6 +7,7 @@ import pytest
 
 from bellows.exception import InvalidCommandError
 from bellows.ezsp import EZSP
+from bellows.ezsp.v9.commands 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