From 54da03b681a1589f836ec3ae10327013a15052bc Mon Sep 17 00:00:00 2001
From: Hernando Castano <hernando@hcastano.com>
Date: Tue, 13 Aug 2024 18:55:59 -0400
Subject: [PATCH] Add `network-jumpstart` subcommand to `entropy-test-cli`

This PR adds a way to trigger a network jumpstart from the test CLI. This is useful for ensuring the
network is in the correct state before registering using the new registration flow.
---
 crates/client/src/client.rs | 28 ++++++++++++++++++++++++++++
 crates/test-cli/src/lib.rs  | 29 +++++++++++++++++++++++++++--
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs
index 4fc87744c..7ca7c9683 100644
--- a/crates/client/src/client.rs
+++ b/crates/client/src/client.rs
@@ -360,3 +360,31 @@ pub async fn change_threshold_accounts(
         .ok_or(anyhow!("Error with transaction"))?;
     Ok(result_event)
 }
+
+/// Trigger a network wide distributed key generation (DKG) event.
+///
+/// Fails if the network has already been jumpstarted.
+pub async fn jumpstart_network(
+    api: &OnlineClient<EntropyConfig>,
+    rpc: &LegacyRpcMethods<EntropyConfig>,
+    signer: sr25519::Pair,
+) -> Result<(), ClientError> {
+    // In this case we don't care too much about the result because we're more interested in the
+    // `FinishedNetworkJumpStart` event, which happens later on.
+    let jump_start_request = entropy::tx().registry().jump_start_network();
+    let _result =
+        submit_transaction_with_pair(api, rpc, &signer, &jump_start_request, None).await?;
+
+    let mut blocks_sub = api.blocks().subscribe_finalized().await?;
+
+    while let Some(block) = blocks_sub.next().await {
+        let block = block?;
+        let events = block.events().await?;
+
+        if events.has::<entropy::registry::events::FinishedNetworkJumpStart>()? {
+            break;
+        }
+    }
+
+    Ok(())
+}
diff --git a/crates/test-cli/src/lib.rs b/crates/test-cli/src/lib.rs
index 2f36ee1de..6af8292da 100644
--- a/crates/test-cli/src/lib.rs
+++ b/crates/test-cli/src/lib.rs
@@ -26,7 +26,7 @@ use entropy_client::{
     },
     client::{
         change_endpoint, change_threshold_accounts, get_accounts, get_api, get_programs, get_rpc,
-        register, sign, store_program, update_programs, VERIFYING_KEY_LENGTH,
+        jumpstart_network, register, sign, store_program, update_programs, VERIFYING_KEY_LENGTH,
     },
 };
 use sp_core::{sr25519, Hasher, Pair};
@@ -76,7 +76,7 @@ enum CliCommand {
         /// A name or mnemonic from which to derive a program modification keypair.
         /// This is used to send the register extrinsic so it must be funded
         /// If giving a name it must be preceded with "//", eg: "--mnemonic-option //Alice"
-        /// If giving a mnemonic it must be enclosed in quotes, eg: "--mnemonic-option "alarm mutual concert...""  
+        /// If giving a mnemonic it must be enclosed in quotes, eg: "--mnemonic-option "alarm mutual concert...""
         #[arg(short, long)]
         mnemonic_option: Option<String>,
     },
@@ -143,6 +143,17 @@ enum CliCommand {
     },
     /// Display a list of registered Entropy accounts
     Status,
+    /// Triggers the network wide distributed key generation process.
+    ///
+    /// A fully jumpstarted network is required for the on-chain registration flow to work
+    /// correctly.
+    ///
+    /// Note: Any account may trigger the jumpstart process.
+    JumpstartNetwork {
+        /// The mnemonic for the signer which will trigger the jumpstart process.
+        #[arg(short, long)]
+        mnemonic_option: Option<String>,
+    },
 }
 
 pub async fn run_command(
@@ -392,6 +403,20 @@ pub async fn run_command(
 
             Ok("Threshold accounts changed".to_string())
         },
+        CliCommand::JumpstartNetwork { mnemonic_option } => {
+            let mnemonic = if let Some(mnemonic_option) = mnemonic_option {
+                mnemonic_option
+            } else {
+                passed_mnemonic.unwrap_or("//Alice".to_string())
+            };
+
+            let signer = <sr25519::Pair as Pair>::from_string(&mnemonic, None)?;
+            println!("Account being used for jumpstart: {}", signer.public());
+
+            jumpstart_network(&api, &rpc, signer).await?;
+
+            Ok("Succesfully jumpstarted network.".to_string())
+        },
     }
 }