Move NVRAM migration into ZNP

puddly · puddly · commit 069abd0e7885 · 2021-11-16T14:08:10.000-05:00
diff --git a/tests/api/test_nvram_migration.py b/tests/api/test_nvram_migration.py
@@ -46,7 +46,7 @@ async def test_addrmgr_empty_entries(make_connected_znp, device):
 
 
 @pytest.mark.parametrize("device", [FormedZStack3CC2531])
-async def test_addrmgr_rewrite_fix(device, make_application, mocker):
+async def test_addrmgr_rewrite_fix(device, make_connected_znp):
     # Keep track of reads
     addrmgr_reads = []
 
@@ -62,7 +62,7 @@ async def test_addrmgr_rewrite_fix(device, make_application, mocker):
         extAddr=t.EUI64.convert("FF:FF:FF:FF:FF:FF:FF:FF"),
     )
 
-    app, znp_server = await make_application(server_cls=device)
+    znp, znp_server = await make_connected_znp(server_cls=device)
     znp_server.callback_for_response(
         c.SYS.OSALNVReadExt.Req(Id=OsalNvIds.ADDRMGR, Offset=0), addrmgr_reads.append
     )
@@ -83,8 +83,7 @@ async def test_addrmgr_rewrite_fix(device, make_application, mocker):
     assert old_addrmgr != nvram[OsalNvIds.ADDRMGR]
 
     assert len(addrmgr_reads) == 0
-    await app.startup()
-    await app.shutdown()
+    await znp.migrate_nvram()
     assert len(addrmgr_reads) == 2
 
     # Bad entries have been fixed
@@ -96,9 +95,7 @@ async def test_addrmgr_rewrite_fix(device, make_application, mocker):
 
     # Will not be read again
     assert len(addrmgr_reads) == 2
-    await app.connect()
-    await app.startup()
-    await app.shutdown()
+    await znp.migrate_nvram()
     assert len(addrmgr_reads) == 2
 
     # Will be migrated again if the migration NVID is deleted
@@ -107,9 +104,7 @@ async def test_addrmgr_rewrite_fix(device, make_application, mocker):
     old_addrmgr2 = nvram[OsalNvIds.ADDRMGR]
 
     assert len(addrmgr_reads) == 2
-    await app.connect()
-    await app.startup()
-    await app.shutdown()
+    await znp.migrate_nvram()
     assert len(addrmgr_reads) == 3
 
     # But nothing will change
diff --git a/zigpy_znp/api.py b/zigpy_znp/api.py
@@ -38,6 +38,8 @@
 STARTUP_TIMEOUT = 15
 NETWORK_COMMISSIONING_TIMEOUT = 30
 
+NVRAM_MIGRATION_ID = 1
+
 
 class ZNP:
     def __init__(self, config: conf.ConfigType):
@@ -457,6 +459,11 @@ async def write_network_info(
             counter_increment=0,
         )
 
+        # Prevent an unnecessary NVRAM migration from running
+        await self.nvram.osal_write(
+            OsalNvIds.ZIGPY_ZNP_MIGRATION_ID, t.uint8_t(NVRAM_MIGRATION_ID), create=True
+        )
+
         if self.version == 1.2:
             await self.nvram.osal_write(
                 OsalNvIds.HAS_CONFIGURED_ZSTACK1,
@@ -472,6 +479,65 @@ async def write_network_info(
 
         LOGGER.debug("Done!")
 
+    async def migrate_nvram(self) -> bool:
+        """
+        Migrates NVRAM entries using the `ZIGPY_ZNP_MIGRATION_ID` NVRAM item.
+        Returns `True` if a migration was performed, `False` otherwise.
+        """
+
+        from zigpy_znp.znp import security
+
+        try:
+            migration_id = await self.nvram.osal_read(
+                OsalNvIds.ZIGPY_ZNP_MIGRATION_ID, item_type=t.uint8_t
+            )
+        except KeyError:
+            migration_id = 0
+
+        initial_migration_id = migration_id
+
+        # Migration 1: empty `ADDRMGR` entries are version-dependent and were improperly
+        #              written for CC253x devices.
+        #
+        #              This migration is stateless and can safely be run more than once:
+        #              the only downside is that startup times increase by 10s on newer
+        #              coordinators, which is why the migration ID is persisted.
+        if migration_id < 1:
+            try:
+                entries = await security.read_addr_manager_entries(self)
+            except KeyError:
+                pass
+            else:
+                fixed_entries = []
+
+                for entry in entries:
+                    if entry.extAddr != t.EUI64.convert("FF:FF:FF:FF:FF:FF:FF:FF"):
+                        fixed_entries.append(entry)
+                    elif self.version == 3.30:
+                        fixed_entries.append(const.EMPTY_ADDR_MGR_ENTRY_ZSTACK3)
+                    else:
+                        fixed_entries.append(const.EMPTY_ADDR_MGR_ENTRY_ZSTACK1)
+
+                if entries != fixed_entries:
+                    LOGGER.warning(
+                        "Repairing %d invalid empty address manager entries (total %d)",
+                        sum(i != j for i, j in zip(entries, fixed_entries)),
+                        len(entries),
+                    )
+                    await security.write_addr_manager_entries(self, fixed_entries)
+
+            migration_id = 1
+
+        if initial_migration_id == migration_id:
+            return False
+
+        await self.nvram.osal_write(
+            OsalNvIds.ZIGPY_ZNP_MIGRATION_ID, t.uint8_t(migration_id), create=True
+        )
+        await self.reset()
+
+        return True
+
     async def reset(self) -> None:
         """
         Performs a soft reset within Z-Stack.
diff --git a/zigpy_znp/zigbee/application.py b/zigpy_znp/zigbee/application.py
@@ -28,7 +28,6 @@
 import zigpy_znp.config as conf
 import zigpy_znp.commands as c
 from zigpy_znp.api import ZNP
-from zigpy_znp.znp import security
 from zigpy_znp.utils import combine_concurrent_calls
 from zigpy_znp.exceptions import CommandNotRecognized, InvalidCommandResponse
 from zigpy_znp.types.nvids import OsalNvIds
@@ -182,7 +181,7 @@ async def write_network_info(
 
     async def start_network(self, *, read_only=False):
         if not read_only:
-            await self._migrate_nvram()
+            await self._znp.migrate_nvram()
             await self._write_stack_settings(reset_if_changed=True)
 
         if self.znp_config[conf.CONF_TX_POWER] is not None:
@@ -779,63 +778,6 @@ async def _set_led_mode(self, *, led: t.uint8_t, mode: c.util.LEDMode) -> None:
         except (asyncio.TimeoutError, CommandNotRecognized):
             LOGGER.info("This build of Z-Stack does not appear to support LED control")
 
-    async def _migrate_nvram(self) -> bool:
-        """
-        Migrates NVRAM entries using the `ZIGPY_ZNP_MIGRATION_ID` NVRAM item.
-        Returns `True` if a migration was performed, `False` otherwise.
-        """
-
-        try:
-            migration_id = await self._znp.nvram.osal_read(
-                OsalNvIds.ZIGPY_ZNP_MIGRATION_ID, item_type=t.uint8_t
-            )
-        except KeyError:
-            migration_id = 0
-
-        initial_migration_id = migration_id
-
-        # Migration 1: empty `ADDRMGR` entries are version-dependent and were improperly
-        #              written for CC253x devices.
-        #
-        #              This migration is stateless and can safely be run more than once:
-        #              the only downside is that startup times increase by 10s on newer
-        #              coordinators, which is why the migration ID is persisted.
-        if migration_id < 1:
-            try:
-                entries = await security.read_addr_manager_entries(self._znp)
-            except KeyError:
-                pass
-            else:
-                fixed_entries = []
-
-                for entry in entries:
-                    if entry.extAddr != t.EUI64.convert("FF:FF:FF:FF:FF:FF:FF:FF"):
-                        fixed_entries.append(entry)
-                    elif self._znp.version == 3.30:
-                        fixed_entries.append(const.EMPTY_ADDR_MGR_ENTRY_ZSTACK3)
-                    else:
-                        fixed_entries.append(const.EMPTY_ADDR_MGR_ENTRY_ZSTACK1)
-
-                if entries != fixed_entries:
-                    LOGGER.warning(
-                        "Repairing %d invalid empty address manager entries (total %d)",
-                        sum(i != j for i, j in zip(entries, fixed_entries)),
-                        len(entries),
-                    )
-                    await security.write_addr_manager_entries(self._znp, fixed_entries)
-
-            migration_id = 1
-
-        if initial_migration_id == migration_id:
-            return False
-
-        await self._znp.nvram.osal_write(
-            OsalNvIds.ZIGPY_ZNP_MIGRATION_ID, t.uint8_t(migration_id), create=True
-        )
-        await self._znp.reset()
-
-        return True
-
     async def _write_stack_settings(self, *, reset_if_changed: bool) -> None:
         """
         Writes network-independent Z-Stack settings to NVRAM.
