Fix unit tests

zigpy · Jun 5, 2024 · 0f6ad34 · 0f6ad34
1 parent bcf2774
commit 0f6ad34
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 18 deletions.
diff --git a/bellows/ezsp/__init__.py b/bellows/ezsp/__init__.py
@@ -128,6 +128,7 @@ async def startup_reset(self) -> None:
             await self.reset()
 
         await self.version()
+        await self.get_xncp_features()
 
     @classmethod
     async def initialize(cls, zigpy_config: dict) -> EZSP:
@@ -180,24 +181,21 @@ async def version(self):
             self._switch_protocol_version(ver)
             await self._command("version", ver)
 
-        try:
-            self._xncp_features = await self.xncp_get_supported_firmware_features()
-        except InvalidCommandError:
-            self._xncp_features = xncp.FirmwareFeatures.NONE
-
         LOGGER.debug(
-            (
-                "EZSP Stack Type: %s"
-                ", Stack Version: %04x"
-                ", Protocol version: %s"
-                ", XNCP features: %s"
-            ),
+            ("EZSP Stack Type: %s" ", Stack Version: %04x" ", Protocol version: %s"),
             stack_type,
             stack_version,
             ver,
-            self._xncp_features,
         )
 
+    async def get_xncp_features(self) -> None:
+        try:
+            self._xncp_features = await self.xncp_get_supported_firmware_features()
+        except InvalidCommandError:
+            self._xncp_features = xncp.FirmwareFeatures.NONE
+
+        LOGGER.debug("XNCP features: %s", self._xncp_features)
+
     def close(self):
         self.stop_ezsp()
         if self._gw:
@@ -426,7 +424,7 @@ async def get_mfg_token(self, token: t.EzspMfgTokenId) -> bytes:
 
     async def _get_mfg_custom_eui_64(self) -> t.EUI64 | None:
         """Get the custom EUI 64 manufacturing token, if it has a valid value."""
-        (data,) = await self.get_mfg_token(t.EzspMfgTokenId.MFG_CUSTOM_EUI_64)
+        data = await self.get_mfg_token(t.EzspMfgTokenId.MFG_CUSTOM_EUI_64)
 
         # Manufacturing tokens do not exist in RCP firmware: all reads are empty
         if not data:

diff --git a/tests/test_application.py b/tests/test_application.py
@@ -76,9 +76,10 @@ def ezsp_mock(ieee):
     )
     mock_ezsp.add_transient_link_key = AsyncMock(return_value=t.EmberStatus.SUCCESS)
     mock_ezsp._protocol = AsyncMock()
-    mock_ezsp.get_supported_firmware_features = AsyncMock(
+    mock_ezsp.xncp_get_supported_firmware_features = AsyncMock(
         return_value=FirmwareFeatures.NONE
     )
+    mock_ezsp._xncp_features = FirmwareFeatures.NONE
 
     type(mock_ezsp).types = ezsp_t7
     type(mock_ezsp).is_ezsp_running = PropertyMock(return_value=True)
@@ -226,9 +227,10 @@ async def nop_mock():
         ]
     )
     ezsp_mock.setMulticastTableEntry = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
-    ezsp_mock.get_supported_firmware_features = AsyncMock(
+    ezsp_mock.xncp_get_supported_firmware_features = AsyncMock(
         return_value=FirmwareFeatures.NONE
     )
+    ezsp_mock._xncp_features = FirmwareFeatures.NONE
 
     app.permit = AsyncMock()
 

diff --git a/tests/test_ezsp.py b/tests/test_ezsp.py
@@ -282,8 +282,9 @@ async def test_no_close_without_callback(ezsp_f):
 
 @patch.object(ezsp.EZSP, "version", new_callable=AsyncMock)
 @patch.object(ezsp.EZSP, "reset", new_callable=AsyncMock)
+@patch.object(ezsp.EZSP, "get_xncp_features", new_callable=AsyncMock)
 @patch("bellows.uart.connect", return_value=MagicMock(spec_set=uart.Gateway))
-async def test_ezsp_init(conn_mock, reset_mock, version_mock):
+async def test_ezsp_init(conn_mock, xncp_mock, reset_mock, version_mock):
     """Test initialize method."""
     zigpy_config = config.CONFIG_SCHEMA({"device": DEVICE_CONFIG})
     await ezsp.EZSP.initialize(zigpy_config)
@@ -643,8 +644,9 @@ async def test_write_custom_eui64_rcp(ezsp_f):
 
 @patch.object(ezsp.EZSP, "version", new_callable=AsyncMock)
 @patch.object(ezsp.EZSP, "reset", new_callable=AsyncMock)
+@patch.object(ezsp.EZSP, "get_xncp_features", new_callable=AsyncMock)
 @patch("bellows.uart.connect", return_value=MagicMock(spec_set=uart.Gateway))
-async def test_ezsp_init_zigbeed(conn_mock, reset_mock, version_mock):
+async def test_ezsp_init_zigbeed(conn_mock, xncp_mock, reset_mock, version_mock):
     """Test initialize method with a received startup reset frame."""
     zigpy_config = config.CONFIG_SCHEMA(
         {
@@ -667,9 +669,12 @@ async def test_ezsp_init_zigbeed(conn_mock, reset_mock, version_mock):
 
 @patch.object(ezsp.EZSP, "version", new_callable=AsyncMock)
 @patch.object(ezsp.EZSP, "reset", new_callable=AsyncMock)
+@patch.object(ezsp.EZSP, "get_xncp_features", new_callable=AsyncMock)
 @patch("bellows.uart.connect", return_value=MagicMock(spec_set=uart.Gateway))
 @patch("bellows.ezsp.NETWORK_COORDINATOR_STARTUP_RESET_WAIT", 0.01)
-async def test_ezsp_init_zigbeed_timeout(conn_mock, reset_mock, version_mock):
+async def test_ezsp_init_zigbeed_timeout(
+    conn_mock, xncp_mock, reset_mock, version_mock
+):
     """Test initialize method with a received startup reset frame."""
     zigpy_config = config.CONFIG_SCHEMA(
         {