sns-testing: make ICP treasury configurable

Currently, the identity that transfer ICPs to the swap participants or
the user provided account is hardcoded. Which limits the usage of the
'sns-testing' to the network created by 'sns-testing-init'.

This commit adds an optional argument to choose the dfx identity to be
used as an ICP treasury to 'sns-testing-init' binary and
'swap-complete' and 'transfer-icp' subcommands of 'sns-testing'.

Author: rvem

rs/sns/testing/src/bin/init.rs

@@ -23,12 +23,18 @@ async fn nns_init(args: NnsInitArgs) {
     let initial_mutations = load_registry_mutations(registry_proto_path);
     let dev_principal_id = get_identity_principal(&args.dev_identity).unwrap();
 
+    let treasury_principal_id = if let Some(icp_treasury_identity) = args.icp_treasury_identity {
+        get_identity_principal(&icp_treasury_identity).unwrap()
+    } else {
+        *TREASURY_PRINCIPAL_ID
+    };
+
     bootstrap_nns(
         &pocket_ic,
         vec![initial_mutations],
         vec![
             (
-                (*TREASURY_PRINCIPAL_ID).into(),
+                treasury_principal_id.into(),
                 Tokens::from_tokens(10_000_000).unwrap(),
             ),
             (dev_principal_id.into(), Tokens::from_tokens(100).unwrap()),

rs/sns/testing/src/bin/main.rs

@@ -87,7 +87,9 @@ async fn validate_network(network: String) {
 }
 
 async fn swap_complete(network: String, args: SwapCompleteArgs) {
-    let agent = get_agent(&network, None).await.unwrap();
+    let agent = get_agent(&network, args.icp_treasury_identity.clone())
+        .await
+        .unwrap();
 
     // Normally, SNSes would have different names, so the vector below would have a single element.
     let target_snses = find_sns_by_name(&agent, args.sns_name.clone()).await;
@@ -113,7 +115,14 @@ async fn swap_complete(network: String, args: SwapCompleteArgs) {
                 };
             neurons_to_follow.extend(principal_neurons);
         }
-        if let Err(e) = complete_sns_swap(&agent, sns.swap, sns.governance, neurons_to_follow).await
+        if let Err(e) = complete_sns_swap(
+            &agent,
+            args.icp_treasury_identity.is_none(),
+            sns.swap,
+            sns.governance,
+            neurons_to_follow,
+        )
+        .await
         {
             eprintln!("Failed to complete swap for SNS {}: {}", args.sns_name, e);
             exit(1);
@@ -154,17 +163,24 @@ async fn sns_proposal_upvote(network: String, args: SnsProposalUpvoteArgs) {
 }
 
 pub async fn transfer_icp(network: String, args: TransferICPArgs) {
-    let agent = get_agent(&network, None).await.unwrap();
+    let agent = get_agent(&network, args.icp_treasury_identity.clone())
+        .await
+        .unwrap();
     let recipient = match args.recipient {
         TransferRecipientArg::Account { account } => account.0,
         TransferRecipientArg::Principal {
             principal_id,
             subaccount,
         } => AccountIdentifier::new(principal_id, subaccount.map(|a| Subaccount(a.0))).to_address(),
     };
-    transfer_icp_from_treasury(&agent, recipient, args.amount.0)
-        .await
-        .unwrap();
+    transfer_icp_from_treasury(
+        &agent,
+        args.icp_treasury_identity.is_none(),
+        recipient,
+        args.amount.0,
+    )
+    .await
+    .unwrap();
 }
 
 #[tokio::main]

rs/sns/testing/src/lib.rs

@@ -41,8 +41,8 @@ pub struct ValidateNetworkArgs {}
 
 #[derive(Debug, Parser)]
 pub struct RunBasicScenarioArgs {
-    /// The name of the 'dfx' identity to use for the scenario. This identity is used to submit NNS
-    /// proposal to create the new SNS and is added as an initial neuron in the new SNS.
+    /// The name of the 'dfx' identity to use for the scenario. The principal of this identity
+    /// is used to submit NNS proposal to create the new SNS and is added as an initial neuron in the new SNS.
     #[arg(long)]
     pub dev_identity: Option<String>,
     /// The ID of the canister to be controlled by the SNS created in the scenario.
@@ -56,6 +56,12 @@ pub struct SwapCompleteArgs {
     /// The name of the SNS to complete the swap for.
     #[arg(long)]
     pub sns_name: String,
+    /// The name of the 'dfx' identity. The principal of this identity is supposed to have
+    /// an account in the ICP ledger with sufficient balance to provide direct participations
+    /// for the SNS swap.
+    /// If not provided, the ephemeral identity with hardcoded principal will be used.
+    #[arg(long)]
+    pub icp_treasury_identity: Option<String>,
     /// The neuron that swap participants will follow.
     #[clap(long, group = "neuron-follow-selection")]
     pub follow_neuron: Option<ParsedSnsNeuron>,
@@ -128,6 +134,11 @@ pub struct TransferICPArgs {
     /// The amount of e8s to transfer.
     #[arg(long)]
     pub amount: ParsedTokens,
+    /// The name of the 'dfx' identity. The principal of this identity is supposed to
+    /// have an account in the ICP ledger with sufficient balance to perform the transfer.
+    /// If not provided, the ephemeral identity with hardcoded principal will be used.
+    #[arg(long)]
+    pub icp_treasury_identity: Option<String>,
 }
 
 /// A wrapper to parse ICP tokens from a decimal number.
@@ -220,4 +231,11 @@ pub struct NnsInitArgs {
     /// hotkey for the NNS neuron with the majority voting power.
     #[arg(long)]
     pub dev_identity: String,
+    /// The name of the 'dfx' identity. The principal of this identity will be added to the
+    /// ICP ledger with 10_000_000 ICP. This identity will be used to transfer ICP
+    /// to provide sufficient direct participations for the SNS swap or to transfer ICP
+    /// to the user provided account.
+    /// If not provided, the ephemeral identity with hardcoded principal will be used.
+    #[arg(long)]
+    pub icp_treasury_identity: Option<String>,
 }

rs/sns/testing/src/sns.rs

@@ -124,6 +124,7 @@ pub async fn create_sns<C: CallCanisters + ProgressNetwork + BuildEphemeralAgent
     let sns_swap = sns.swap;
     complete_sns_swap(
         neuron_agent,
+        true,
         sns_swap,
         sns_governance,
         if follow_dev_neuron {
@@ -186,16 +187,24 @@ async fn get_current_participation<C: CallCanisters>(
 
 // Completes the swap by transferring the required amount of ICP from the "treasury" account
 // and participating in the swap for each participant using agents provided as arguments:
-// 1) swap_treasury_agent - Agent for the identity that has sufficient amout of ICP tokens to close the swap.
-// 2) swap_participants_agents - Agents for the identities that will participate in the swap.
-// 3) swap_canister - SNS Swap canister IDs.
+// 1) agent - Agent that is used to provide IC network settings.
+// 2) use_ephemeral_icp_treasury - defines whether the identity of the 'agent' or 'TREASURY_PRINCIPAL_ID'.
+//    is used to transfer ICP to the swap participants.
+// 2) swap_canister - SNS Swap canister ID.
+// 3) governance_canister - SNS Governance canister ID.
+// 4) neurons_to_follow - SNS Neuron IDs that will be followed by the swap participants.
 pub async fn complete_sns_swap<C: CallCanisters + ProgressNetwork + BuildEphemeralAgent>(
     agent: &C,
+    use_ephemeral_icp_treasury: bool,
     swap_canister: SwapCanister,
     governance_canister: GovernanceCanister,
     neurons_to_follow: Vec<SnsNeuronId>,
 ) -> Result<(), String> {
-    let swap_treasury_agent = &agent.build_ephemeral_agent(TREASURY_SECRET_KEY.clone());
+    let swap_treasury_agent = if use_ephemeral_icp_treasury {
+        &agent.build_ephemeral_agent(TREASURY_SECRET_KEY.clone())
+    } else {
+        agent
+    };
 
     println!("Waiting for the swap to be open...");
     await_swap_lifecycle(agent, swap_canister, Lifecycle::Open, true).await?;
@@ -514,7 +523,6 @@ pub async fn propose_sns_controlled_test_canister_upgrade<C: CallCanisters + Pro
     .unwrap()
 }
 
-
 // Waits for the upgrade proposal to be adopted and executed and then waits for the canister to become available
 // after upgrade using arguments:
 // 1) agent - Agent that will be used to check the status of the canister.

rs/sns/testing/src/utils.rs

@@ -252,12 +252,17 @@ pub async fn validate_target_canister<C: CallCanisters>(
 
 pub async fn transfer_icp_from_treasury<C: CallCanisters + BuildEphemeralAgent>(
     agent: &C,
+    use_ephemeral_icp_treasury: bool,
     to: AccountIdBlob,
     amount: Tokens,
 ) -> Result<u64, TransferError> {
-    let treasury_agent = agent.build_ephemeral_agent(TREASURY_SECRET_KEY.clone());
+    let treasury_agent = if use_ephemeral_icp_treasury {
+        &agent.build_ephemeral_agent(TREASURY_SECRET_KEY.clone())
+    } else {
+        agent
+    };
     transfer(
-        &treasury_agent,
+        treasury_agent,
         TransferArgs {
             to,
             amount,