From 7f26d8218b8926762645a335986e55bd9040b09d Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Fri, 30 Jun 2023 14:40:55 -0400
Subject: [PATCH] Handle `getTokenData` command missing (#571)

---
 bellows/exception.py     |  4 ++++
 bellows/ezsp/__init__.py | 15 +++++++--------
 bellows/ezsp/protocol.py |  4 ++--
 tests/test_ezsp.py       |  6 +++++-
 4 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/bellows/exception.py b/bellows/exception.py
index 25fbb311..bbc66203 100644
--- a/bellows/exception.py
+++ b/bellows/exception.py
@@ -5,5 +5,9 @@ class EzspError(APIException):
     pass
 
 
+class InvalidCommandError(EzspError):
+    pass
+
+
 class ControllerError(ControllerException):
     pass
diff --git a/bellows/ezsp/__init__.py b/bellows/ezsp/__init__.py
index 4b53fdad..220f22c4 100644
--- a/bellows/ezsp/__init__.py
+++ b/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)
diff --git a/bellows/ezsp/protocol.py b/bellows/ezsp/protocol.py
index 87fcb3e7..1a1aa055 100644
--- a/bellows/ezsp/protocol.py
+++ b/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}"
                         )
diff --git a/tests/test_ezsp.py b/tests/test_ezsp.py
index 22820c96..795876d4 100644
--- a/tests/test_ezsp.py
+++ b/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."""