E2E tests for coldkey swap

.github/workflows/typescript-e2e.yml

@@ -103,6 +103,8 @@ jobs:
             binary: release
           - test: zombienet_staking
             binary: fast
+          - test: zombienet_coldkey_swap
+            binary: fast
 
     name: "typescript-e2e-${{ matrix.test }}"
 

ts-tests/moonwall.config.json

@@ -94,6 +94,32 @@
                     "descriptor": "subtensor"
                 }
             ]
+        },
+        {
+            "name": "zombienet_coldkey_swap",
+            "timeout": 600000,
+            "testFileDir": ["suites/zombienet_coldkey_swap"],
+            "runScripts": [
+                "generate-types.sh",
+                "build-spec.sh"
+            ],
+            "foundation": {
+                "type": "zombie",
+                "zombieSpec": {
+                    "configPath": "./configs/zombie_node.json",
+                    "skipBlockCheck": []
+                }
+            },
+            "vitestArgs": {
+                "bail": 1
+            },
+            "connections": [
+                {
+                    "name": "Node",
+                    "type": "papi",
+                    "endpoints": ["ws://127.0.0.1:9947"]
+                }
+            ]
         }, {
             "name": "smoke_mainnet",
             "testFileDir": ["suites/smoke"],

ts-tests/scripts/generate-types.sh

@@ -58,4 +58,4 @@ if [ "$GENERATE_TYPES" = true ]; then
   exit 0
 else
   echo "==> Types are up-to-date, nothing to do."
-fi
+fi

ts-tests/suites/zombienet_coldkey_swap/01-coldkey-swap-sudo.test.ts

@@ -0,0 +1,175 @@
+import { expect, beforeAll } from "vitest";
+import { describeSuite } from "@moonwall/cli";
+import { subtensor } from "@polkadot-api/descriptors";
+import type { TypedApi } from "polkadot-api";
+import {
+    addNewSubnetwork,
+    addStake,
+    announceColdkeySwap,
+    burnedRegister,
+    coldkeyHashBinary,
+    disputeColdkeySwap,
+    forceSetBalance,
+    generateKeyringPair,
+    getBalance,
+    getColdkeySwapAnnouncement,
+    getColdkeySwapDispute,
+    getHotkeyOwner,
+    getOwnedHotkeys,
+    getStake,
+    getStakingHotkeys,
+    getSubnetOwner,
+    startCall,
+    sudoResetColdkeySwap,
+    sudoSwapColdkey,
+    tao,
+} from "../../utils";
+
+describeSuite({
+    id: "01_coldkey_swap_sudo",
+    title: "▶ coldkey swap sudo operations",
+    foundationMethods: "zombie",
+    testCases: ({ it, context, log }) => {
+        let api: TypedApi<typeof subtensor>;
+
+        beforeAll(async () => {
+            api = context.papi("Node").getTypedApi(subtensor);
+        });
+
+        it({
+            id: "T01",
+            title: "reset as root: clears announcement and dispute",
+            test: async () => {
+                const oldColdkey = generateKeyringPair("sr25519");
+                const newColdkey = generateKeyringPair("sr25519");
+                await forceSetBalance(api, oldColdkey.address);
+
+                // Announce and dispute
+                await announceColdkeySwap(api, oldColdkey, coldkeyHashBinary(newColdkey));
+                await disputeColdkeySwap(api, oldColdkey);
+                log("Announced + disputed");
+
+                // Verify storage before reset
+                const annBefore = await getColdkeySwapAnnouncement(api, oldColdkey.address);
+                expect(annBefore, "announcement should exist before reset").not.toBeNull();
+
+                // Reset via sudo
+                await sudoResetColdkeySwap(api, oldColdkey.address);
+                log("Reset via sudo");
+
+                // Verify storage cleared
+                const annAfter = await getColdkeySwapAnnouncement(api, oldColdkey.address);
+                expect(annAfter, "announcement should be cleared").toBeNull();
+                const dispAfter = await getColdkeySwapDispute(api, oldColdkey.address);
+                expect(dispAfter, "dispute should be cleared").toBeNull();
+                log("Storage cleared");
+
+                // Re-announce should succeed
+                await announceColdkeySwap(api, oldColdkey, coldkeyHashBinary(newColdkey));
+                log("Re-announce after reset succeeded");
+            },
+        });
+
+        it({
+            id: "T02",
+            title: "instant swap as root: transfers stake and ownership across multiple subnets",
+            test: async () => {
+                const oldColdkey = generateKeyringPair("sr25519");
+                const newColdkey = generateKeyringPair("sr25519");
+                const hotkey1 = generateKeyringPair("sr25519");
+                const hotkey2 = generateKeyringPair("sr25519");
+
+                await forceSetBalance(api, oldColdkey.address);
+                await forceSetBalance(api, hotkey1.address);
+                await forceSetBalance(api, hotkey2.address);
+
+                // Create two subnets
+                const netuid1 = await addNewSubnetwork(api, hotkey1, oldColdkey);
+                await startCall(api, netuid1, oldColdkey);
+                const netuid2 = await addNewSubnetwork(api, hotkey2, oldColdkey);
+                await startCall(api, netuid2, oldColdkey);
+                log(`Created subnets ${netuid1} and ${netuid2}`);
+
+                // Register hotkey1 on subnet2 and stake on both
+                await burnedRegister(api, netuid2, hotkey1.address, oldColdkey);
+                await addStake(api, oldColdkey, hotkey1.address, netuid1, tao(100));
+                await addStake(api, oldColdkey, hotkey1.address, netuid2, tao(50));
+
+                const stake1Before = await getStake(api, hotkey1.address, oldColdkey.address, netuid1);
+                const stake2Before = await getStake(api, hotkey1.address, oldColdkey.address, netuid2);
+                expect(stake1Before, "should have stake on subnet1").toBeGreaterThan(0n);
+                expect(stake2Before, "should have stake on subnet2").toBeGreaterThan(0n);
+                log(`Before — subnet1 stake: ${stake1Before}, subnet2 stake: ${stake2Before}`);
+
+                // Sudo swap
+                await sudoSwapColdkey(api, oldColdkey.address, newColdkey.address, 0n);
+                log("Sudo swap executed");
+
+                // Verify both subnets' stake migrated
+                expect(
+                    await getStake(api, hotkey1.address, oldColdkey.address, netuid1),
+                    "old coldkey stake on subnet1 should be 0"
+                ).toBe(0n);
+                expect(
+                    await getStake(api, hotkey1.address, newColdkey.address, netuid1),
+                    "new coldkey should have stake on subnet1"
+                ).toBeGreaterThan(0n);
+                expect(
+                    await getStake(api, hotkey1.address, oldColdkey.address, netuid2),
+                    "old coldkey stake on subnet2 should be 0"
+                ).toBe(0n);
+                expect(
+                    await getStake(api, hotkey1.address, newColdkey.address, netuid2),
+                    "new coldkey should have stake on subnet2"
+                ).toBeGreaterThan(0n);
+                log("Stake migrated on both subnets");
+
+                // Verify subnet ownership transferred
+                expect(await getSubnetOwner(api, netuid1), "new coldkey should own subnet1").toBe(newColdkey.address);
+                expect(await getSubnetOwner(api, netuid2), "new coldkey should own subnet2").toBe(newColdkey.address);
+
+                // Verify hotkey ownership transferred
+                expect(await getHotkeyOwner(api, hotkey1.address), "hotkey1 owner").toBe(newColdkey.address);
+                expect(await getHotkeyOwner(api, hotkey2.address), "hotkey2 owner").toBe(newColdkey.address);
+
+                // Verify old coldkey is fully empty
+                expect(
+                    (await getOwnedHotkeys(api, oldColdkey.address)).length,
+                    "old coldkey should own no hotkeys"
+                ).toBe(0);
+                expect(
+                    (await getStakingHotkeys(api, oldColdkey.address)).length,
+                    "old coldkey should have no staking hotkeys"
+                ).toBe(0);
+                expect(await getBalance(api, oldColdkey.address), "old coldkey balance should be 0").toBe(0n);
+
+                log("All state migrated across both subnets");
+            },
+        });
+
+        it({
+            id: "T03",
+            title: "instant swap as root: clears pending announcement",
+            test: async () => {
+                const oldColdkey = generateKeyringPair("sr25519");
+                const newColdkey = generateKeyringPair("sr25519");
+                const decoy = generateKeyringPair("sr25519");
+                await forceSetBalance(api, oldColdkey.address);
+
+                // Announce for decoy
+                await announceColdkeySwap(api, oldColdkey, coldkeyHashBinary(decoy));
+                const annBefore = await getColdkeySwapAnnouncement(api, oldColdkey.address);
+                expect(annBefore, "announcement should exist").not.toBeNull();
+                log("Pending announcement exists");
+
+                // Sudo swap with different key
+                await sudoSwapColdkey(api, oldColdkey.address, newColdkey.address, 0n);
+
+                // Announcement should be cleared
+                const annAfter = await getColdkeySwapAnnouncement(api, oldColdkey.address);
+                expect(annAfter, "announcement should be cleared after root swap").toBeNull();
+                log("Announcement cleared by root swap");
+            },
+        });
+    },
+});

ts-tests/utils/coldkey_swap.ts

@@ -0,0 +1,144 @@
+import { Keyring } from "@polkadot/keyring";
+import { blake2AsHex } from "@polkadot/util-crypto";
+import type { KeyringPair } from "@moonwall/util";
+import { waitForTransactionWithRetry } from "./transactions.js";
+import type { TypedApi } from "polkadot-api";
+import type { subtensor } from "@polkadot-api/descriptors";
+import { FixedSizeBinary } from "polkadot-api";
+
+export const ANNOUNCEMENT_DELAY = 10;
+export const REANNOUNCEMENT_DELAY = 10;
+
+/** Compute BLAKE2-256 hash of a keypair's public key as a FixedSizeBinary (used for announcements). */
+export function coldkeyHashBinary(pair: KeyringPair): FixedSizeBinary<32> {
+    return FixedSizeBinary.fromHex(blake2AsHex(pair.publicKey, 256));
+}
+
+/** Compute BLAKE2-256 hash of a keypair's public key as hex string. */
+export function coldkeyHash(pair: KeyringPair): string {
+    return blake2AsHex(pair.publicKey, 256);
+}
+
+// ── Sudo configuration ──────────────────────────────────────────────────
+
+export async function sudoSetAnnouncementDelay(api: TypedApi<typeof subtensor>, delay: number): Promise<void> {
+    const keyring = new Keyring({ type: "sr25519" });
+    const alice = keyring.addFromUri("//Alice");
+    const internalCall = api.tx.AdminUtils.sudo_set_coldkey_swap_announcement_delay({
+        duration: delay,
+    });
+    const tx = api.tx.Sudo.sudo({ call: internalCall.decodedCall });
+    await waitForTransactionWithRetry(api, tx, alice, "sudo_set_coldkey_swap_announcement_delay");
+}
+
+export async function sudoSetReannouncementDelay(api: TypedApi<typeof subtensor>, delay: number): Promise<void> {
+    const keyring = new Keyring({ type: "sr25519" });
+    const alice = keyring.addFromUri("//Alice");
+    const internalCall = api.tx.AdminUtils.sudo_set_coldkey_swap_reannouncement_delay({
+        duration: delay,
+    });
+    const tx = api.tx.Sudo.sudo({ call: internalCall.decodedCall });
+    await waitForTransactionWithRetry(api, tx, alice, "sudo_set_coldkey_swap_reannouncement_delay");
+}
+
+// ── Transaction wrappers (throw on failure) ─────────────────────────────
+
+export async function announceColdkeySwap(
+    api: TypedApi<typeof subtensor>,
+    signer: KeyringPair,
+    newColdkeyHash: FixedSizeBinary<32>
+): Promise<void> {
+    const tx = api.tx.SubtensorModule.announce_coldkey_swap({
+        new_coldkey_hash: newColdkeyHash,
+    });
+    await waitForTransactionWithRetry(api, tx, signer, "announce_coldkey_swap");
+}
+
+export async function swapColdkeyAnnounced(
+    api: TypedApi<typeof subtensor>,
+    signer: KeyringPair,
+    newAddress: string
+): Promise<void> {
+    const tx = api.tx.SubtensorModule.swap_coldkey_announced({
+        new_coldkey: newAddress,
+    });
+    await waitForTransactionWithRetry(api, tx, signer, "swap_coldkey_announced");
+}
+
+export async function disputeColdkeySwap(api: TypedApi<typeof subtensor>, signer: KeyringPair): Promise<void> {
+    const tx = api.tx.SubtensorModule.dispute_coldkey_swap();
+    await waitForTransactionWithRetry(api, tx, signer, "dispute_coldkey_swap");
+}
+
+export async function clearColdkeySwapAnnouncement(
+    api: TypedApi<typeof subtensor>,
+    signer: KeyringPair
+): Promise<void> {
+    const tx = api.tx.SubtensorModule.clear_coldkey_swap_announcement();
+    await waitForTransactionWithRetry(api, tx, signer, "clear_coldkey_swap_announcement");
+}
+
+export async function sudoResetColdkeySwap(api: TypedApi<typeof subtensor>, coldkeyAddress: string): Promise<void> {
+    const keyring = new Keyring({ type: "sr25519" });
+    const alice = keyring.addFromUri("//Alice");
+    const internalCall = api.tx.SubtensorModule.reset_coldkey_swap({
+        coldkey: coldkeyAddress,
+    });
+    const tx = api.tx.Sudo.sudo({ call: internalCall.decodedCall });
+    await waitForTransactionWithRetry(api, tx, alice, "sudo_reset_coldkey_swap");
+}
+
+export async function sudoSwapColdkey(
+    api: TypedApi<typeof subtensor>,
+    oldAddress: string,
+    newAddress: string,
+    swapCost = 0n
+): Promise<void> {
+    const keyring = new Keyring({ type: "sr25519" });
+    const alice = keyring.addFromUri("//Alice");
+    const internalCall = api.tx.SubtensorModule.swap_coldkey({
+        old_coldkey: oldAddress,
+        new_coldkey: newAddress,
+        swap_cost: swapCost,
+    });
+    const tx = api.tx.Sudo.sudo({ call: internalCall.decodedCall });
+    await waitForTransactionWithRetry(api, tx, alice, "sudo_swap_coldkey");
+}
+
+// ── Storage query helpers ───────────────────────────────────────────────
+
+export async function getColdkeySwapAnnouncement(
+    api: TypedApi<typeof subtensor>,
+    address: string
+): Promise<{ when: number; hash: string } | null> {
+    const result = await api.query.SubtensorModule.ColdkeySwapAnnouncements.getValue(address);
+    if (!result) return null;
+    const [when, hash] = result;
+    return { when, hash: hash.asHex() };
+}
+
+export async function getColdkeySwapDispute(api: TypedApi<typeof subtensor>, address: string): Promise<number | null> {
+    const result = await api.query.SubtensorModule.ColdkeySwapDisputes.getValue(address);
+    if (result === undefined) return null;
+    return Number(result);
+}
+
+/** Get the owner coldkey of a hotkey. */
+export async function getHotkeyOwner(api: TypedApi<typeof subtensor>, hotkey: string): Promise<string> {
+    return await api.query.SubtensorModule.Owner.getValue(hotkey);
+}
+
+/** Get the list of hotkeys owned by a coldkey. */
+export async function getOwnedHotkeys(api: TypedApi<typeof subtensor>, coldkey: string): Promise<string[]> {
+    return await api.query.SubtensorModule.OwnedHotkeys.getValue(coldkey);
+}
+
+/** Get the list of hotkeys a coldkey is staking to. */
+export async function getStakingHotkeys(api: TypedApi<typeof subtensor>, coldkey: string): Promise<string[]> {
+    return await api.query.SubtensorModule.StakingHotkeys.getValue(coldkey);
+}
+
+/** Get the owner coldkey of a subnet. */
+export async function getSubnetOwner(api: TypedApi<typeof subtensor>, netuid: number): Promise<string> {
+    return await api.query.SubtensorModule.SubnetOwner.getValue(netuid);
+}

ts-tests/utils/index.ts

@@ -4,3 +4,4 @@ export * from "./subnet.js";
 export * from "./staking.js";
 export * from "./shield_helpers.ts";
 export * from "./account.ts";
+export * from "./coldkey_swap.ts";