Implement can_write_network_settings API (#687)

* Implement `can_write_network_settings` API

* Fix unit tests

* Bump minimum zigpy version

* Modify whitespace to trigger GitHub actions cache bust

* Fix unit tests to account for new zigpy packet priority

* Add some tests

Author: puddly

bellows/zigbee/application.py

@@ -19,7 +19,11 @@
 import zigpy.config
 import zigpy.device
 import zigpy.endpoint
-from zigpy.exceptions import NetworkNotFormed
+from zigpy.exceptions import (
+    CannotWriteNetworkSettings,
+    DestructiveWriteNetworkSettings,
+    NetworkNotFormed,
+)
 import zigpy.state
 import zigpy.types
 import zigpy.util
@@ -346,49 +350,77 @@ async def load_network_info(self, *, load_devices=False) -> None:
         async for nwk, eui64 in ezsp.read_address_table():
             self.state.network_info.nwk_addresses[eui64] = nwk
 
+    async def can_write_network_settings(
+        self,
+        *,
+        network_info: zigpy.state.NetworkInfo,
+        node_info: zigpy.state.NodeInfo,
+    ) -> bool:
+        # If we are not overwriting the EUI64, we can always write the network info
+        if node_info.ieee == zigpy.types.EUI64.UNKNOWN:
+            return True
+
+        ezsp = self._ezsp
+        stack_specific = network_info.stack_specific.get("ezsp", {})
+        (current_eui64,) = await ezsp.getEui64()
+
+        # If we are not actually replacing the EUI64, we can write the network info
+        if node_info.ieee == current_eui64:
+            return True
+
+        # If the adapter supports EUI64 rewriting, we can always write the network info
+        if await ezsp.can_rewrite_custom_eui64():
+            return True
+
+        # If the adapter does not support EUI64 rewriting, we cannot proceed
+        if not await ezsp.can_burn_userdata_custom_eui64():
+            _, _, version = await self._get_board_info()
+            raise CannotWriteNetworkSettings(
+                f"Please upgrade your adapter firmware. The adapter IEEE address"
+                f" has been overwritten and firmware {version!r} does not support"
+                f" writing it a second time."
+            )
+
+        # Otherwise, we need explicit confirmation
+        if not stack_specific.get(
+            "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it"
+        ):
+            _, _, version = await self._get_board_info()
+            raise DestructiveWriteNetworkSettings(
+                f"Please upgrade your adapter firmware. The adapter IEEE address"
+                f" needs to be replaced and firmware {version!r} does not support"
+                f" writing it multiple times."
+            )
+
+        return True
+
     async def write_network_info(
         self, *, network_info: zigpy.state.NetworkInfo, node_info: zigpy.state.NodeInfo
     ) -> None:
         ezsp = self._ezsp
 
+        # Throw an error early if we can't actually restore
+        await self.can_write_network_settings(
+            network_info=network_info, node_info=node_info
+        )
         await self.reset_network_info()
 
         stack_specific = network_info.stack_specific.get("ezsp", {})
         (current_eui64,) = await ezsp.getEui64()
-        wrote_eui64 = False
 
         if (
             node_info.ieee != zigpy.types.EUI64.UNKNOWN
             and node_info.ieee != current_eui64
         ):
             if await ezsp.can_rewrite_custom_eui64():
                 await ezsp.write_custom_eui64(node_info.ieee)
-                wrote_eui64 = True
-            elif not stack_specific.get(
-                "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it"
-            ):
-                _, _, version = await self._get_board_info()
-                raise ControllerError(
-                    f"Please upgrade your adapter firmware. The adapter IEEE address"
-                    f" needs to be replaced and firmware {version!r} does not support"
-                    f" writing it multiple times."
-                )
-            elif not await ezsp.can_burn_userdata_custom_eui64():
-                _, _, version = await self._get_board_info()
-                raise ControllerError(
-                    f"Please upgrade your adapter firmware. The adapter IEEE address"
-                    f" has been overwritten and firmware {version!r} does not support"
-                    f" writing it a second time."
-                )
             else:
+                assert await ezsp.can_burn_userdata_custom_eui64()
                 await ezsp.write_custom_eui64(node_info.ieee, burn_into_userdata=True)
-                wrote_eui64 = True
 
-        if wrote_eui64:
-            # Reset after writing the EUI64, as it touches NVRAM
             await self._reset()
         else:
-            # If we cannot write the new EUI64, don't mess up key entries with the
+            # If we do not write the new EUI64, don't mess up key entries with the
             # unwritten EUI64 address
             node_info.ieee = current_eui64
             network_info.tc_link_key.partner_ieee = current_eui64

pyproject.toml

@@ -17,7 +17,7 @@ dependencies = [
     "click",
     "click-log>=0.2.1",
     "voluptuous",
-    "zigpy>=0.79.0",
+    "zigpy>=0.83.0",
     'async-timeout; python_version<"3.11"',
 ]
 

requirements_test.txt

@@ -17,4 +17,4 @@ pytest-asyncio>=0.17
 pytest>=7.1.3
 zigpy>=0.54.1
 ruff==0.0.261
-Flake8-pyproject
+Flake8-pyproject

tests/test_application.py

@@ -756,7 +756,7 @@ def send_unicast(aps_frame, data, message_tag, nwk):
     )
 
     await app.send_packet(packet)
-    app._concurrent_requests_semaphore.max_value = 10000
+    app._concurrent_requests_semaphore.max_concurrency = 10000
     results = await asyncio.gather(
         *(app.send_packet(packet) for _ in range(256 + 1)), return_exceptions=True
     )
@@ -1024,7 +1024,7 @@ async def test_send_packet_unicast_concurrency(app, packet, monkeypatch):
     monkeypatch.setattr(bellows.zigbee.application, "MESSAGE_SEND_TIMEOUT_MAINS", 0.5)
     monkeypatch.setattr(bellows.zigbee.application, "MESSAGE_SEND_TIMEOUT_BATTERY", 0.5)
 
-    app._concurrent_requests_semaphore.max_value = 10
+    app._concurrent_requests_semaphore.max_concurrency = 12
 
     max_concurrency = 0
     in_flight_requests = 0
@@ -1073,9 +1073,14 @@ async def send_unicast(nwk, aps_frame, message_tag, data):
 
     app._ezsp.send_unicast = AsyncMock(side_effect=send_unicast)
 
-    responses = await asyncio.gather(*[app.send_packet(packet) for _ in range(100)])
+    responses = await asyncio.gather(
+        *[
+            app.send_packet(packet.replace(priority=zigpy_t.PacketPriority.HIGH))
+            for _ in range(100)
+        ]
+    )
     assert len(responses) == 100
-    assert max_concurrency == 10
+    assert max_concurrency == 12
     assert in_flight_requests == 0
 
 
@@ -2027,6 +2032,96 @@ async def test_write_network_info(
     ]
 
 
+@pytest.mark.parametrize(
+    ("node_ieee", "current_eui64", "can_rewrite", "can_burn", "confirmation_flag"),
+    [
+        (
+            zigpy_t.EUI64.UNKNOWN,
+            t.EUI64.convert("00:01:02:03:04:05:06:07"),
+            False,
+            False,
+            False,
+        ),
+        (
+            t.EUI64.convert("00:01:02:03:04:05:06:07"),
+            t.EUI64.convert("00:01:02:03:04:05:06:07"),
+            False,
+            False,
+            False,
+        ),
+        (
+            t.EUI64.convert("aa:aa:aa:aa:aa:aa:aa:aa"),
+            t.EUI64.convert("00:01:02:03:04:05:06:07"),
+            True,
+            False,
+            False,
+        ),
+        (
+            t.EUI64.convert("aa:aa:aa:aa:aa:aa:aa:aa"),
+            t.EUI64.convert("00:01:02:03:04:05:06:07"),
+            False,
+            True,
+            True,
+        ),
+    ],
+)
+async def test_can_write_network_settings(
+    app: ControllerApplication,
+    zigpy_backup: zigpy.backups.NetworkBackup,
+    node_ieee: t.EUI64,
+    current_eui64: t.EUI64,
+    can_rewrite: bool,
+    can_burn: bool,
+    confirmation_flag: bool,
+) -> None:
+    """Test `can_write_network_settings`."""
+    app._ezsp.getEui64 = AsyncMock(return_value=[current_eui64])
+    app._ezsp.can_rewrite_custom_eui64 = AsyncMock(return_value=can_rewrite)
+    app._ezsp.can_burn_userdata_custom_eui64 = AsyncMock(return_value=can_burn)
+    app._get_board_info = AsyncMock(
+        return_value=("Mock board", "Mock Manufacturer", "Mock version")
+    )
+
+    node_info = zigpy_backup.node_info.replace(ieee=node_ieee)
+    network_info = zigpy_backup.network_info
+
+    if confirmation_flag:
+        network_info = network_info.replace(
+            stack_specific={
+                "ezsp": {
+                    **network_info.stack_specific.get("ezsp", {}),
+                    "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it": True,
+                }
+            }
+        )
+
+    assert await app.can_write_network_settings(
+        network_info=network_info,
+        node_info=node_info,
+    )
+
+
+async def test_write_network_info_uses_write_custom_eui64(
+    app: ControllerApplication,
+    zigpy_backup: zigpy.backups.NetworkBackup,
+) -> None:
+    """Test that `write_network_info` uses `write_custom_eui64` correctly."""
+    app._ezsp.can_rewrite_custom_eui64 = AsyncMock(return_value=True)
+    app._ezsp.write_custom_eui64 = AsyncMock()
+
+    different_ieee = t.EUI64.convert("aa:aa:aa:aa:aa:aa:aa:aa")
+    node_info = zigpy_backup.node_info.replace(ieee=different_ieee)
+
+    with patch.object(app, "_reset"):
+        await app.write_network_info(
+            node_info=node_info,
+            network_info=zigpy_backup.network_info,
+        )
+
+    # Verify write_custom_eui64 was called without burn_into_userdata flag
+    assert app._ezsp.write_custom_eui64.mock_calls == [call(different_ieee)]
+
+
 async def test_network_scan(app: ControllerApplication) -> None:
     app._ezsp._protocol.startScan.return_value = [t.sl_Status.OK]
 
@@ -2209,7 +2304,7 @@ async def test_migration_failure_eui64_overwrite_confirmation(
     # Migration explicitly fails if we need to write the EUI64 but the adapter treats it
     # as a write-once operation
     with pytest.raises(
-        zigpy.exceptions.ControllerException,
+        zigpy.exceptions.DestructiveWriteNetworkSettings,
         match=(
             "Please upgrade your adapter firmware. The adapter IEEE address needs to be"
             " replaced and firmware 'Mock version' does not support writing it multiple"
@@ -2229,7 +2324,7 @@ async def test_migration_failure_eui64_overwrite_confirmation(
         app._ezsp, "write_custom_eui64", wraps=app._ezsp.write_custom_eui64
     ), patch.object(app._ezsp, "can_burn_userdata_custom_eui64", return_value=False):
         with pytest.raises(
-            zigpy.exceptions.ControllerException,
+            zigpy.exceptions.CannotWriteNetworkSettings,
             match=(
                 "Please upgrade your adapter firmware. The adapter IEEE address has"
                 " been overwritten and firmware 'Mock version' does not support writing"