Update write_ieee to use the new bellows API (#87)

* Update `write_ieee` to use the new bellows API

* Try probing EZSP first

* Adjust logging

* Add a unit test

* Adjust README to include `--force` flag

* Re-run pre-commit

Author: puddly

README.md

@@ -74,3 +74,6 @@ $ universal-silabs-flasher \
 The IEEE address can also be specified without colons: `--ieee 003c84fffe92bb2c`.
 
 If the current device's IEEE address already matches the provided one, the command will not write it unnecessarily.
+Depending on firmware version, writing the IEEE address can be a **permanent** operation. If this is the case,
+you will need to upgrade the firmware on your adapter to a more recent release of EmberZNet or perform the one-time
+write with `--force`.

tests/test_flasher.py

@@ -0,0 +1,39 @@
+import asyncio
+from unittest.mock import call, patch
+
+import zigpy.types as t
+
+from universal_silabs_flasher.common import Version
+from universal_silabs_flasher.flasher import Flasher, ProbeResult
+
+
+async def test_write_emberznet_eui64():
+    flasher = Flasher(device="/dev/ttyMOCK")
+
+    with (
+        patch.object(
+            flasher, "probe_gecko_bootloader", side_effect=asyncio.TimeoutError
+        ),
+        patch.object(
+            flasher,
+            "probe_ezsp",
+            return_value=ProbeResult(
+                version=Version("7.4.4.0 build 0"),
+                continue_probing=False,
+                baudrate=115200,
+            ),
+        ),
+        patch.object(flasher, "_connect_ezsp") as mock_connect_ezsp,
+    ):
+        ezsp = mock_connect_ezsp.return_value.__aenter__.return_value
+
+        ezsp.getEui64.return_value = (t.EUI64.convert("00:11:22:33:44:55:66:77"),)
+        ezsp.write_custom_eui64.return_value = None
+
+        await flasher.write_emberznet_eui64(
+            new_ieee=t.EUI64.convert("11:22:33:44:55:66:77:88"), force=True
+        )
+
+    assert ezsp.write_custom_eui64.mock_calls == [
+        call(ieee=t.EUI64.convert("11:22:33:44:55:66:77:88"), burn_into_userdata=True)
+    ]

universal_silabs_flasher/flash.py

@@ -11,7 +11,6 @@
 import typing
 import urllib.parse
 
-import bellows.types
 import click
 import coloredlogs
 import zigpy.ota.validators
@@ -248,12 +247,11 @@ async def probe(ctx: click.Context) -> None:
 @main.command()
 @click.pass_context
 @click.option("--ieee", required=True, type=zigpy.types.EUI64.convert)
+@click.option("--force", default=False, type=bool)
 @click_coroutine
-async def write_ieee(ctx: click.Context, ieee: zigpy.types.EUI64) -> None:
-    new_eui64 = bellows.types.EmberEUI64(ieee)
-
+async def write_ieee(ctx: click.Context, ieee: zigpy.types.EUI64, force: bool) -> None:
     try:
-        await ctx.obj["flasher"].write_emberznet_eui64(new_eui64)
+        await ctx.obj["flasher"].write_emberznet_eui64(ieee, force=force)
     except (ValueError, RuntimeError) as e:
         raise click.ClickException(str(e)) from e
 

universal_silabs_flasher/flasher.py

@@ -9,6 +9,7 @@
 import bellows.ezsp
 import bellows.types
 from zigpy.serial import SerialProtocol
+import zigpy.types
 
 from .common import (
     PROBE_TIMEOUT,
@@ -162,10 +163,18 @@ async def probe_spinel(self, baudrate: int) -> ProbeResult:
     async def probe_app_type(
         self,
         types: typing.Iterable[ApplicationType] | None = None,
+        try_first: tuple[ApplicationType, ...] = (),
     ) -> None:
         if types is None:
             types = self._probe_methods
 
+        # fmt: off
+        types = (
+              [m for m in types if m in try_first]
+            + [m for m in types if m not in try_first]
+        )
+        # fmt: on
+
         # Reset into bootloader
         if self._reset_target:
             await self.enter_bootloader_reset(self._reset_target)
@@ -205,6 +214,8 @@ async def probe_app_type(
             except asyncio.TimeoutError:
                 continue
 
+            _LOGGER.debug("Probe result: %s", result)
+
             # Keep track of the bootloader version for later
             if probe_method == ApplicationType.GECKO_BOOTLOADER:
                 _LOGGER.info("Detected bootloader version %s", result.version)
@@ -307,30 +318,25 @@ async def dump_emberznet_config(self) -> None:
                     continue
                 print(f"{config.name}={v[1]}")
 
-    async def write_emberznet_eui64(self, new_eui64: bellows.types.EUI64) -> bool:
-        await self.probe_app_type()
+    async def write_emberznet_eui64(
+        self, new_ieee: zigpy.types.EUI64, force: bool = False
+    ) -> bool:
+        await self.probe_app_type(
+            try_first=[ApplicationType.GECKO_BOOTLOADER, ApplicationType.EZSP]
+        )
 
         if self.app_type != ApplicationType.EZSP:
             raise RuntimeError(f"Device is not running EmberZNet: {self.app_type}")
 
         async with self._connect_ezsp(self.app_baudrate) as ezsp:
-            (current_eui64,) = await ezsp.getEui64()
-            _LOGGER.info("Current device IEEE: %s", current_eui64)
+            (current_ieee,) = await ezsp.getEui64()
+            _LOGGER.info("Current device IEEE: %s", current_ieee)
 
-            if current_eui64 == new_eui64:
+            if current_ieee == new_ieee:
                 _LOGGER.info("Device IEEE address already matches, not overwriting")
                 return False
 
-            if not await ezsp.can_write_custom_eui64():
-                raise ValueError(
-                    "IEEE address has already been written, it cannot be written again"
-                )
-
-            (status,) = await ezsp.setMfgToken(
-                bellows.types.EzspMfgTokenId.MFG_CUSTOM_EUI_64, new_eui64.serialize()
-            )
-
-            if status != bellows.types.EmberStatus.SUCCESS:
-                raise RuntimeError(f"Failed to write IEEE address: {status}")
+            await ezsp.write_custom_eui64(ieee=new_ieee, burn_into_userdata=force)
+            _LOGGER.info("Wrote new device IEEE: %s", new_ieee)
 
         return True