Correctly write the NWK update ID when forming a network

bellows/ezsp/__init__.py

@@ -515,6 +515,40 @@ async def write_custom_eui64(
                 f" cannot be written again without erasing flash."
             )
 
+    async def write_nwk_update_id(self, nwk_update_id: int) -> None:
+        """Write NWK update ID to NVRAM token.
+
+        This is a workaround for the lack of EZSP API to set the network update ID.
+        The EmberNetworkParameters.nwkUpdateId field is ignored during formNetwork(),
+        so we must write directly to the NVRAM token.
+        """
+        try:
+            # Read current network management token
+            rsp = await self.getTokenData(
+                token=t.NV3KeyId.NVM3KEY_STACK_NETWORK_MANAGEMENT, index=0
+            )
+            assert t.sl_Status.from_ember_status(rsp.status) == t.sl_Status.OK
+        except (InvalidCommandError, AttributeError, AssertionError):
+            LOGGER.warning("NV3 interface not available, cannot write NWK update ID")
+            return
+
+        # Deserialize current token
+        token, remaining = t.NV3StackNetworkManagementToken.deserialize(rsp.value)
+        assert not remaining
+
+        # Update the NWK update ID
+        updated_token = token.replace(update_id=t.uint8_t(nwk_update_id))
+
+        # Write updated token back to NVRAM
+        (status,) = await self.setTokenData(
+            token=t.NV3KeyId.NVM3KEY_STACK_NETWORK_MANAGEMENT,
+            index=0,
+            token_data=updated_token.serialize(),
+        )
+        assert t.sl_Status.from_ember_status(status) == t.sl_Status.OK
+
+        LOGGER.debug("Updated NWK update ID to %d in NVRAM", nwk_update_id)
+
     def add_callback(self, cb):
         id_ = hash(cb)
         while id_ in self._callbacks:

bellows/types/struct.py

@@ -706,3 +706,12 @@ class EmberMultiPhyRadioParameters(EzspStruct):
     radioTxPower: basic.int8s
     radioPage: basic.uint8_t
     radioChannel: basic.uint8_t
+
+
+class NV3StackNetworkManagementToken(EzspStruct):
+    """NV3 stack network management token value."""
+
+    active_channels: named.Channels
+    manager_node_id: named.NWK
+    update_id: basic.uint8_t
+    padding: basic.uint8_t

bellows/zigbee/application.py

@@ -433,6 +433,12 @@ async def write_network_info(
         parameters.channels = t.Channels(network_info.channel_mask)
 
         await ezsp.formNetwork(parameters=parameters)
+
+        # Write NWK update ID to NVRAM after network formation. This is needed because
+        # formNetwork() appears to ignore or reset the nwkUpdateId field
+        if network_info.nwk_update_id != 0:
+            await ezsp.write_nwk_update_id(network_info.nwk_update_id)
+
         await self._ensure_network_running()
 
     async def reset_network_info(self):

tests/test_ezsp.py

@@ -644,6 +644,86 @@ async def test_write_custom_eui64_rcp(ezsp_f):
     ]
 
 
+async def test_write_nwk_update_id(ezsp_f):
+    """Test writing network update ID to NVRAM token."""
+    # Mock the token data response
+    mock_token_data = t.NV3StackNetworkManagementToken(
+        active_channels=t.Channels(0x07FFF800),
+        manager_node_id=t.NWK(0x0000),
+        update_id=t.uint8_t(0),
+        padding=t.uint8_t(0),
+    )
+
+    ezsp_f.getTokenData = AsyncMock(
+        return_value=GetTokenDataRsp(
+            status=t.EmberStatus.SUCCESS,
+            value=mock_token_data.serialize(),
+        )
+    )
+    ezsp_f.setTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
+
+    # Test writing update ID
+    await ezsp_f.write_nwk_update_id(7)
+
+    # Verify getTokenData was called
+    ezsp_f.getTokenData.assert_called_once_with(
+        token=t.NV3KeyId.NVM3KEY_STACK_NETWORK_MANAGEMENT,
+        index=0,
+    )
+
+    # Verify setTokenData was called with updated token
+    ezsp_f.setTokenData.assert_called_once()
+    call_args = ezsp_f.setTokenData.call_args[1]
+    assert call_args["token"] == t.NV3KeyId.NVM3KEY_STACK_NETWORK_MANAGEMENT
+    assert call_args["index"] == 0
+
+    # Deserialize the token data to verify the update ID was set correctly
+    updated_token, remaining = t.NV3StackNetworkManagementToken.deserialize(
+        call_args["token_data"]
+    )
+    assert not remaining
+    assert updated_token.update_id == 7
+    assert updated_token.active_channels == t.Channels(0x07FFF800)
+    assert updated_token.manager_node_id == t.NWK(0x0000)
+
+
+async def test_write_nwk_update_id_nv3_unavailable(ezsp_f):
+    """Test writing network update ID when NV3 is not available."""
+    # Mock InvalidCommandError to simulate NV3 not being available
+    ezsp_f.getTokenData = AsyncMock(
+        side_effect=InvalidCommandError("NV3 not available")
+    )
+
+    # Should not raise an exception, just log a warning
+    await ezsp_f.write_nwk_update_id(7)
+
+    # Verify getTokenData was called
+    ezsp_f.getTokenData.assert_called_once_with(
+        token=t.NV3KeyId.NVM3KEY_STACK_NETWORK_MANAGEMENT,
+        index=0,
+    )
+
+
+async def test_write_nwk_update_id_failure(ezsp_f):
+    """Test writing network update ID when token read fails."""
+    # Mock failure response
+    ezsp_f.getTokenData = AsyncMock(
+        return_value=GetTokenDataRsp(
+            status=t.EmberStatus.INVALID_CALL,
+            value=b"",
+        )
+    )
+
+    # Should not raise an exception, just log a warning
+    await ezsp_f.write_nwk_update_id(7)
+
+    # Verify getTokenData was called
+    ezsp_f.getTokenData.assert_called_once_with(
+        token=t.NV3KeyId.NVM3KEY_STACK_NETWORK_MANAGEMENT,
+        index=0,
+    )
+
+
 @patch.object(EZSP, "version", new_callable=AsyncMock)
 @patch.object(EZSP, "reset", new_callable=AsyncMock)
 @patch.object(EZSP, "get_xncp_features", new_callable=AsyncMock)