Handle getTokenData command missing (#571)

bellows/exception.py

@@ -5,5 +5,9 @@ class EzspError(APIException):
     pass
 
 
+class InvalidCommandError(EzspError):
+    pass
+
+
 class ControllerError(ControllerException):
     pass

bellows/ezsp/__init__.py

@@ -19,7 +19,7 @@
 import zigpy.config
 
 import bellows.config as conf
-from bellows.exception import EzspError
+from bellows.exception import EzspError, InvalidCommandError
 import bellows.types as t
 import bellows.uart
 
@@ -374,17 +374,16 @@ async def get_board_info(self) -> tuple[str, str, str]:
 
     async def _get_nv3_restored_eui64_key(self) -> t.NV3KeyId | None:
         """Get the NV3 key for the device's restored EUI64, if one exists."""
-        try:
-            # If the EZSP version doesn't have `getTokenData`, it doesn't implement NV3
-            self.getTokenData
-        except AttributeError:
-            return None
-
         for key in (
             t.NV3KeyId.CREATOR_STACK_RESTORED_EUI64,  # NCP firmware
             t.NV3KeyId.NVM3KEY_STACK_RESTORED_EUI64,  # RCP firmware
         ):
-            status, data = await self.getTokenData(key, 0)
+            try:
+                status, data = 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.EmberEUI64.deserialize(data)

bellows/ezsp/protocol.py

@@ -13,7 +13,7 @@
     from asyncio import timeout as asyncio_timeout  # pragma: no cover
 
 from bellows.config import CONF_EZSP_CONFIG, CONF_EZSP_POLICIES
-from bellows.exception import EzspError
+from bellows.exception import InvalidCommandError
 from bellows.typing import GatewayType
 
 LOGGER = logging.getLogger(__name__)
@@ -192,7 +192,7 @@ def __call__(self, data: bytes) -> None:
                 if frame_name == "invalidCommand":
                     sent_cmd_name = self.COMMANDS_BY_ID[expected_id][0]
                     future.set_exception(
-                        EzspError(
+                        InvalidCommandError(
                             f"{sent_cmd_name} command is an {frame_name}, was sent "
                             f"under {sequence} sequence number: {result[0].name}"
                         )

tests/test_ezsp.py

@@ -504,7 +504,11 @@ async def _mock_cmd(*args, **kwargs):
 
 @pytest.mark.parametrize(
     "value, expected_result",
-    [(b"\xFF" * 8, True), (bytes.fromhex("0846b8a11c004b1200"), False), (b"", False)],
+    [
+        (b"\xFF" * 8, True),
+        (bytes.fromhex("0846b8a11c004b1200"), False),
+        (b"", False),
+    ],
 )
 async def test_can_burn_userdata_custom_eui64(ezsp_f, value, expected_result):
     """Test detecting if a custom EUI64 has been written."""