feat(ICP-Rosetta): FI-1540: add disburse of neuron functionality (#2182)

This MR proposes the following changes:

1. Add the functionality of disbursing a neuron to the ICP Rosetta
client

---------

Co-authored-by: Mathias Björkqvist <mathias.bjorkqvist@dfinity.org>

rs/rosetta-api/icp/BUILD.bazel

@@ -79,6 +79,7 @@ DEV_DEPENDENCIES = [
     "//rs/nns/governance/init",
     "//rs/nns/handlers/root/impl:root",
     "//rs/nns/test_utils",
+    "//rs/registry/canister",
     "//rs/rosetta-api/icp:rosetta-api",
     "//rs/rosetta-api/icp/client:ic-icp-rosetta-client",
     "//rs/rosetta-api/icp/ledger_canister_blocks_synchronizer/test_utils",
@@ -168,9 +169,11 @@ rust_test_suite_with_extra_srcs(
         "//rs/canister_sandbox",
         "//rs/canister_sandbox:sandbox_launcher",
         "//rs/ledger_suite/icp/ledger:ledger-canister-wasm-notify-method",
-        "//rs/nns/governance:governance-canister",
+        "//rs/nns/governance:governance-canister-test",
+        "//rs/nns/handlers/lifeline/impl:lifeline_canister",
         "//rs/nns/handlers/root/impl:root-canister",
         "//rs/pocket_ic_server:pocket-ic-server",
+        "//rs/registry/canister:registry-canister",
         "//rs/replica",
         "//rs/rosetta-api/icp:ic-rosetta-api-rosetta-blocks",
         "//rs/rosetta-api/icp:rosetta-api",
@@ -184,8 +187,10 @@ rust_test_suite_with_extra_srcs(
         "ROSETTA_BIN_PATH": "$(rootpath //rs/rosetta-api/icp:ic-rosetta-api-rosetta-blocks)",
         "SANDBOX_LAUNCHER": "$(rootpath //rs/canister_sandbox:sandbox_launcher)",
         "ICP_LEDGER_DEPLOYED_VERSION_WASM_PATH": "$(rootpath @mainnet_icp_ledger_canister//file)",
-        "GOVERNANCE_CANISTER_WASM_PATH": "$(rootpath //rs/nns/governance:governance-canister)",
+        "GOVERNANCE_CANISTER_WASM_PATH": "$(rootpath //rs/nns/governance:governance-canister-test)",
         "ROOT_CANISTER_WASM_PATH": "$(rootpath //rs/nns/handlers/root/impl:root-canister)",
+        "REGISTRY_CANISTER_WASM_PATH": "$(rootpath //rs/registry/canister:registry-canister)",
+        "LIFELINE_CANISTER_WASM_PATH": "$(rootpath //rs/nns/handlers/lifeline/impl:lifeline_canister)",
     },
     extra_srcs = glob([
         "tests/system_tests/common/*.rs",

rs/rosetta-api/icp/Cargo.toml

@@ -38,6 +38,7 @@ num-bigint = { workspace = true }
 on_wire = { path = "../../rust_canisters/on_wire" }
 prometheus = { workspace = true }
 rand = { workspace = true }
+registry-canister = { path = "../../registry/canister" }
 reqwest = { workspace = true }
 rolling-file = { workspace = true }
 rosetta-core = { path = "../common/rosetta_core" }

rs/rosetta-api/icp/client/src/lib.rs

@@ -11,6 +11,7 @@ use ic_rosetta_api::models::ConstructionMetadataRequestOptions;
 use ic_rosetta_api::models::ConstructionPayloadsRequestMetadata;
 use ic_rosetta_api::models::OperationIdentifier;
 use ic_rosetta_api::request_types::ChangeAutoStakeMaturityMetadata;
+use ic_rosetta_api::request_types::DisburseMetadata;
 use ic_rosetta_api::request_types::NeuronIdentifierMetadata;
 use ic_rosetta_api::request_types::RequestType;
 use ic_rosetta_api::request_types::SetDissolveTimestampMetadata;
@@ -327,6 +328,35 @@ impl RosettaClient {
         }])
     }
 
+    pub async fn build_disburse_neuron_operations(
+        signer_principal: Principal,
+        neuron_index: u64,
+        recipient: Option<AccountIdentifier>,
+    ) -> anyhow::Result<Vec<Operation>> {
+        Ok(vec![Operation {
+            operation_identifier: OperationIdentifier {
+                index: 0,
+                network_index: None,
+            },
+            related_operations: None,
+            type_: "DISBURSE".to_string(),
+            status: None,
+            account: Some(rosetta_core::identifiers::AccountIdentifier::from(
+                AccountIdentifier::new(PrincipalId(signer_principal), None),
+            )),
+            amount: None,
+            coin_change: None,
+            metadata: Some(
+                DisburseMetadata {
+                    neuron_index,
+                    recipient,
+                }
+                .try_into()
+                .map_err(|e| anyhow::anyhow!("Failed to convert metadata: {:?}", e))?,
+            ),
+        }])
+    }
+
     pub async fn network_list(&self) -> anyhow::Result<NetworkListResponse> {
         self.call_endpoint("/network/list", &MetadataRequest { metadata: None })
             .await
@@ -899,6 +929,33 @@ impl RosettaClient {
         )
         .await
     }
+
+    /// If a neuron is in the state DISSOLVED you can disburse the neuron with this function.
+    pub async fn disburse_neuron<T>(
+        &self,
+        network_identifier: NetworkIdentifier,
+        signer_keypair: &T,
+        disburse_neuron_args: RosettaDisburseNeuronArgs,
+    ) -> anyhow::Result<ConstructionSubmitResponse>
+    where
+        T: RosettaSupportedKeyPair,
+    {
+        let disburse_neuron_operations = RosettaClient::build_disburse_neuron_operations(
+            signer_keypair.generate_principal_id()?.0,
+            disburse_neuron_args.neuron_index,
+            disburse_neuron_args.recipient,
+        )
+        .await?;
+
+        self.make_submit_and_wait_for_transaction(
+            signer_keypair,
+            network_identifier,
+            disburse_neuron_operations,
+            None,
+            None,
+        )
+        .await
+    }
 }
 
 pub struct RosettaTransferArgs {
@@ -1083,6 +1140,41 @@ impl RosettaIncreaseNeuronStakeArgs {
     }
 }
 
+pub struct RosettaIncreaseNeuronStakeArgsBuilder {
+    additional_stake: Nat,
+    neuron_index: Option<u64>,
+    // The subaccount from which the ICP should be transferred
+    from_subaccount: Option<[u8; 32]>,
+}
+
+impl RosettaIncreaseNeuronStakeArgsBuilder {
+    pub fn new(additional_stake: Nat) -> Self {
+        Self {
+            additional_stake,
+            neuron_index: None,
+            from_subaccount: None,
+        }
+    }
+
+    pub fn with_neuron_index(mut self, neuron_index: u64) -> Self {
+        self.neuron_index = Some(neuron_index);
+        self
+    }
+
+    pub fn with_from_subaccount(mut self, from_subaccount: Subaccount) -> Self {
+        self.from_subaccount = Some(from_subaccount);
+        self
+    }
+
+    pub fn build(self) -> RosettaIncreaseNeuronStakeArgs {
+        RosettaIncreaseNeuronStakeArgs {
+            additional_stake: self.additional_stake,
+            neuron_index: self.neuron_index,
+            from_subaccount: self.from_subaccount,
+        }
+    }
+}
+
 pub struct RosettaChangeAutoStakeMaturityArgs {
     pub neuron_index: Option<u64>,
     pub requested_setting_for_auto_stake_maturity: bool,
@@ -1122,38 +1214,39 @@ impl RosettaChangeAutoStakeMaturityArgsBuilder {
         }
     }
 }
+pub struct RosettaDisburseNeuronArgs {
+    pub neuron_index: u64,
+    pub recipient: Option<AccountIdentifier>,
+}
 
-pub struct RosettaIncreaseNeuronStakeArgsBuilder {
-    additional_stake: Nat,
-    neuron_index: Option<u64>,
-    // The subaccount from which the ICP should be transferred
-    from_subaccount: Option<[u8; 32]>,
+impl RosettaDisburseNeuronArgs {
+    pub fn builder(neuron_index: u64) -> RosettaDisburseNeuronArgsBuilder {
+        RosettaDisburseNeuronArgsBuilder::new(neuron_index)
+    }
 }
 
-impl RosettaIncreaseNeuronStakeArgsBuilder {
-    pub fn new(additional_stake: Nat) -> Self {
+pub struct RosettaDisburseNeuronArgsBuilder {
+    neuron_index: u64,
+    recipient: Option<AccountIdentifier>,
+}
+
+impl RosettaDisburseNeuronArgsBuilder {
+    pub fn new(neuron_index: u64) -> Self {
         Self {
-            additional_stake,
-            neuron_index: None,
-            from_subaccount: None,
+            neuron_index,
+            recipient: None,
         }
     }
 
-    pub fn with_neuron_index(mut self, neuron_index: u64) -> Self {
-        self.neuron_index = Some(neuron_index);
+    pub fn with_recipient(mut self, recipient: AccountIdentifier) -> Self {
+        self.recipient = Some(recipient);
         self
     }
 
-    pub fn with_from_subaccount(mut self, from_subaccount: Subaccount) -> Self {
-        self.from_subaccount = Some(from_subaccount);
-        self
-    }
-
-    pub fn build(self) -> RosettaIncreaseNeuronStakeArgs {
-        RosettaIncreaseNeuronStakeArgs {
-            additional_stake: self.additional_stake,
+    pub fn build(self) -> RosettaDisburseNeuronArgs {
+        RosettaDisburseNeuronArgs {
             neuron_index: self.neuron_index,
-            from_subaccount: self.from_subaccount,
+            recipient: self.recipient,
         }
     }
 }

rs/rosetta-api/icp/tests/system_tests/common/system_test_environment.rs

@@ -1,5 +1,6 @@
 use crate::common::utils::get_custom_agent;
-use crate::common::utils::wait_for_rosetta_to_sync_up_to_block;
+use crate::common::utils::get_test_agent;
+use crate::common::utils::wait_for_rosetta_to_catch_up_with_icp_ledger;
 use crate::common::{
     constants::{DEFAULT_INITIAL_BALANCE, STARTING_CYCLES_PER_CANISTER},
     utils::test_identity,
@@ -16,12 +17,16 @@ use ic_icrc1_test_utils::LedgerEndpointArg;
 use ic_icrc1_tokens_u256::U256;
 use ic_ledger_test_utils::build_ledger_wasm;
 use ic_ledger_test_utils::pocket_ic_helpers::ledger::LEDGER_CANISTER_ID;
+use ic_nns_common::init::LifelineCanisterInitPayloadBuilder;
 use ic_nns_constants::GOVERNANCE_CANISTER_ID;
 use ic_nns_constants::LIFELINE_CANISTER_ID;
+use ic_nns_constants::REGISTRY_CANISTER_ID;
 use ic_nns_constants::ROOT_CANISTER_ID;
 use ic_nns_governance_init::GovernanceCanisterInitPayloadBuilder;
 use ic_nns_handler_root::init::RootCanisterInitPayloadBuilder;
 use ic_nns_test_utils::common::build_governance_wasm;
+use ic_nns_test_utils::common::build_lifeline_wasm;
+use ic_nns_test_utils::common::build_registry_wasm;
 use ic_nns_test_utils::common::build_root_wasm;
 use ic_rosetta_test_utils::path_from_env;
 use ic_types::PrincipalId;
@@ -32,6 +37,7 @@ use num_traits::cast::ToPrimitive;
 use pocket_ic::CanisterSettings;
 use pocket_ic::{nonblocking::PocketIc, PocketIcBuilder};
 use prost::Message;
+use registry_canister::init::RegistryCanisterInitPayloadBuilder;
 use rosetta_core::identifiers::NetworkIdentifier;
 use std::collections::HashMap;
 use tempfile::TempDir;
@@ -132,14 +138,6 @@ impl RosettaTestingEnvironment {
     }
 
     pub async fn restart_rosetta_node(mut self, options: RosettaOptions) -> Self {
-        let ledger_tip = self
-            .rosetta_client
-            .network_status(self.network_identifier.clone())
-            .await
-            .unwrap()
-            .current_block_identifier
-            .index;
-
         self.rosetta_context.kill_rosetta_process();
 
         let rosetta_bin = path_from_env("ROSETTA_BIN_PATH");
@@ -149,13 +147,12 @@ impl RosettaTestingEnvironment {
         self.rosetta_client =
             RosettaClient::from_str_url(&format!("http://localhost:{}", self.rosetta_context.port))
                 .expect("Unable to parse url");
-        wait_for_rosetta_to_sync_up_to_block(
+        wait_for_rosetta_to_catch_up_with_icp_ledger(
             &self.rosetta_client,
             self.network_identifier.clone(),
-            ledger_tip,
+            &get_test_agent(self.pocket_ic.url().unwrap().port().unwrap()).await,
         )
-        .await
-        .unwrap();
+        .await;
         self
     }
 }
@@ -280,7 +277,9 @@ impl RosettaTestingEnvironmentBuilder {
                     Some(nns_root_canister_controller),
                 )
                 .await;
-
+            pocket_ic
+                .add_cycles(nns_root_canister_id, STARTING_CYCLES_PER_CANISTER)
+                .await;
             let governance_canister_wasm = build_governance_wasm();
             let governance_canister_id = Principal::from(GOVERNANCE_CANISTER_ID);
             let governance_canister_controller = ROOT_CANISTER_ID.get().0;
@@ -316,6 +315,61 @@ impl RosettaTestingEnvironmentBuilder {
                 .advance_time(std::time::Duration::from_secs(60))
                 .await;
             pocket_ic.tick().await;
+
+            let nns_lifeline_canister_wasm = build_lifeline_wasm();
+            let nns_lifeline_canister_id = Principal::from(LIFELINE_CANISTER_ID);
+            let nns_lifeline_canister_controller = ROOT_CANISTER_ID.get().0;
+            let nns_lifeline_canister = pocket_ic
+                .create_canister_with_id(
+                    Some(nns_lifeline_canister_controller),
+                    Some(CanisterSettings {
+                        controllers: Some(vec![nns_lifeline_canister_controller]),
+                        ..Default::default()
+                    }),
+                    nns_lifeline_canister_id,
+                )
+                .await
+                .expect("Unable to create the NNS Lifeline canister");
+
+            pocket_ic
+                .install_canister(
+                    nns_lifeline_canister,
+                    nns_lifeline_canister_wasm.bytes().to_vec(),
+                    Encode!(&LifelineCanisterInitPayloadBuilder::new().build()).unwrap(),
+                    Some(nns_lifeline_canister_controller),
+                )
+                .await;
+            pocket_ic
+                .add_cycles(nns_lifeline_canister_id, STARTING_CYCLES_PER_CANISTER)
+                .await;
+
+            let nns_registry_canister_wasm = build_registry_wasm();
+            let nns_registry_canister_id = Principal::from(REGISTRY_CANISTER_ID);
+            let nns_registry_canister_controller = ROOT_CANISTER_ID.get().0;
+            let nns_registry_canister = pocket_ic
+                .create_canister_with_id(
+                    Some(nns_registry_canister_controller),
+                    Some(CanisterSettings {
+                        controllers: Some(vec![nns_registry_canister_controller]),
+                        ..Default::default()
+                    }),
+                    nns_registry_canister_id,
+                )
+                .await
+                .expect("Unable to create the NNS Registry canister");
+
+            pocket_ic
+                .install_canister(
+                    nns_registry_canister,
+                    nns_registry_canister_wasm.bytes().to_vec(),
+                    Encode!(&RegistryCanisterInitPayloadBuilder::new().build()).unwrap(),
+                    Some(nns_registry_canister_controller),
+                )
+                .await;
+
+            pocket_ic
+                .add_cycles(nns_registry_canister_id, STARTING_CYCLES_PER_CANISTER)
+                .await;
         }
 
         let replica_url = pocket_ic.make_live(None).await;
@@ -384,19 +438,12 @@ impl RosettaTestingEnvironmentBuilder {
             .unwrap();
 
         // Wait for rosetta to catch up with the ledger
-        if let Some(last_block_idx) = block_idxes.last() {
-            let rosetta_last_block_idx = wait_for_rosetta_to_sync_up_to_block(
-                &rosetta_client,
-                network_identifier.clone(),
-                *last_block_idx,
-            )
-            .await;
-            assert_eq!(
-                Some(*last_block_idx),
-                rosetta_last_block_idx,
-                "Wait for rosetta sync failed."
-            );
-        }
+        wait_for_rosetta_to_catch_up_with_icp_ledger(
+            &rosetta_client,
+            network_identifier.clone(),
+            &get_test_agent(replica_port).await,
+        )
+        .await;
 
         RosettaTestingEnvironment {
             pocket_ic,