Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

test(nns/sns): Introduce NNS suite builder for PocketIc #3666

Merged
merged 2 commits into from
Jan 31, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 0 additions & 14 deletions rs/nervous_system/agent/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,20 +27,6 @@ pub trait Request: Send {
type Response: CandidType + DeserializeOwned;
}

// impl<R: ic_nervous_system_clients::Request> Request for R {
// fn method(&self) -> &'static str {
// Self::METHOD
// }
// fn update(&self) -> bool {
// Self::UPDATE
// }
// fn payload(&self) -> Result<Vec<u8>, candid::Error> {
// candid::encode_one(self)
// }

// type Response = <Self as ic_nervous_system_clients::Request>::Response;
// }

aterga marked this conversation as resolved.
Show resolved Hide resolved
pub trait CallCanisters: sealed::Sealed {
type Error: Display + Send + std::error::Error + 'static;
fn call<R: Request>(
Expand Down
322 changes: 213 additions & 109 deletions rs/nervous_system/integration_tests/src/pocket_ic_helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -314,6 +314,209 @@ pub async fn add_wasms_to_sns_wasm(
})
}

/// A builder for fine-tuning and installing the NNS canister suite in PocketIc.
pub struct NnsInstaller {
aterga marked this conversation as resolved.
Show resolved Hide resolved
mainnet_nns_canister_versions: Option<bool>,
neurons_fund_hotkeys: Vec<PrincipalId>,
custom_initial_registry_mutations: Option<Vec<RegistryAtomicMutateRequest>>,
aterga marked this conversation as resolved.
Show resolved Hide resolved
initial_balances: Vec<(AccountIdentifier, Tokens)>,
}

impl Default for NnsInstaller {
fn default() -> Self {
Self::new()
}
}

impl NnsInstaller {
fn new() -> Self {
Self {
// Enforce an explicit decision.
aterga marked this conversation as resolved.
Show resolved Hide resolved
mainnet_nns_canister_versions: None,
neurons_fund_hotkeys: vec![],
custom_initial_registry_mutations: None,
initial_balances: vec![],
}
}

/// Requests the mainnet Wasm versions for all NNS canisters being installed.
pub fn with_mainnet_nns_canister_versions(&mut self) -> &mut Self {
self.mainnet_nns_canister_versions = Some(true);
self
}

/// Requests tip-of-this-branch Wasm versions for all NNS canisters being installed.
pub fn with_tip_nns_canister_versions(&mut self) -> &mut Self {
aterga marked this conversation as resolved.
Show resolved Hide resolved
self.mainnet_nns_canister_versions = Some(false);
self
}

/// Requests that the NNS Governance is initialized with a Neurons' Fund neuron controlled
/// by `neurons_fund_hotkeys`.
aterga marked this conversation as resolved.
Show resolved Hide resolved
pub fn with_neurons_fund_hotkeys(
&mut self,
neurons_fund_hotkeys: Vec<PrincipalId>,
) -> &mut Self {
self.neurons_fund_hotkeys = neurons_fund_hotkeys;
self
}

/// Requests that the NNS Registry is initialized with `custom_initial_registry_mutations`.
///
/// The custom mutations must result in an invariant-compliant Registry state.
///
/// Without this specification, default initial Registry mutations will be used, which are
/// guaranteed to be compliant.
pub fn with_custom_initial_registry_mutations(
&mut self,
custom_initial_registry_mutations: Vec<RegistryAtomicMutateRequest>,
) -> &mut Self {
self.custom_initial_registry_mutations = Some(custom_initial_registry_mutations);
self
}

/// Requests `initial_balances` as a Vec of `(test_user_icp_ledger_account,
/// test_user_icp_ledger_initial_balance)` pairs, representing some initial ICP balances.
pub fn with_initial_balances(
&mut self,
initial_balances: Vec<(AccountIdentifier, Tokens)>,
) -> &mut Self {
self.initial_balances = initial_balances;
self
}

/// Installs the NNS canister suite.
///
/// Ensures that there is a whale neuron with `TEST_NEURON_1_ID`.
///
/// Requires that the PocketIC instance has an NNS and an SNS subnet.
///
/// Returns the list of `controller_principal_id`s of pre-configured NNS neurons.
pub async fn install(self, pocket_ic: &PocketIc) -> Vec<PrincipalId> {
let with_mainnet_canister_versions = self
.mainnet_nns_canister_versions
.expect("Please explicitly request either mainnet or tip-of-the-branch NNS version.");

let topology = pocket_ic.topology().await;

let sns_subnet_id = topology.get_sns().expect("No SNS subnet found");
let sns_subnet_id = SubnetId::from(PrincipalId::from(sns_subnet_id));

let mut nns_init_payload_builder = NnsInitPayloadsBuilder::new();

if let Some(custom_initial_registry_mutations) = self.custom_initial_registry_mutations {
nns_init_payload_builder.with_initial_mutations(custom_initial_registry_mutations);
} else {
nns_init_payload_builder.with_initial_invariant_compliant_mutations();
}

let maturity_equivalent_icp_e8s = 1_500_000 * E8;
nns_init_payload_builder
.with_test_neurons_fund_neurons_with_hotkeys(
self.neurons_fund_hotkeys,
maturity_equivalent_icp_e8s,
)
.with_sns_dedicated_subnets(vec![sns_subnet_id])
.with_sns_wasm_access_controls(true);

for (test_user_icp_ledger_account, test_user_icp_ledger_initial_balance) in
self.initial_balances
{
nns_init_payload_builder.with_ledger_account(
test_user_icp_ledger_account,
test_user_icp_ledger_initial_balance,
);
}

let nns_init_payload = nns_init_payload_builder.build();

let (governance_wasm, ledger_wasm, root_wasm, lifeline_wasm, sns_wasm_wasm, registry_wasm) =
if with_mainnet_canister_versions {
(
build_mainnet_governance_wasm(),
build_mainnet_ledger_wasm(),
build_mainnet_root_wasm(),
build_mainnet_lifeline_wasm(),
build_mainnet_sns_wasms_wasm(),
build_mainnet_registry_wasm(),
)
} else {
(
build_governance_wasm(),
build_ledger_wasm(),
build_root_wasm(),
build_lifeline_wasm(),
build_sns_wasms_wasm(),
build_registry_wasm(),
)
};

install_canister(
pocket_ic,
"ICP Ledger",
LEDGER_CANISTER_ID,
Encode!(&nns_init_payload.ledger).unwrap(),
ledger_wasm,
Some(ROOT_CANISTER_ID.get()),
)
.await;
install_canister(
pocket_ic,
"NNS Root",
ROOT_CANISTER_ID,
Encode!(&nns_init_payload.root).unwrap(),
root_wasm,
Some(LIFELINE_CANISTER_ID.get()),
)
.await;
install_canister(
pocket_ic,
"NNS Governance",
GOVERNANCE_CANISTER_ID,
nns_init_payload.governance.encode_to_vec(),
governance_wasm,
Some(ROOT_CANISTER_ID.get()),
)
.await;
install_canister(
pocket_ic,
"Lifeline",
LIFELINE_CANISTER_ID,
Encode!(&nns_init_payload.lifeline).unwrap(),
lifeline_wasm,
Some(ROOT_CANISTER_ID.get()),
)
.await;
install_canister(
pocket_ic,
"NNS SNS-W",
SNS_WASM_CANISTER_ID,
Encode!(&nns_init_payload.sns_wasms).unwrap(),
sns_wasm_wasm,
Some(ROOT_CANISTER_ID.get()),
)
.await;
install_canister(
pocket_ic,
"Registry",
REGISTRY_CANISTER_ID,
Encode!(&nns_init_payload.registry).unwrap(),
registry_wasm,
Some(ROOT_CANISTER_ID.get()),
)
.await;

let nns_neurons = nns_init_payload
.governance
.neurons
.values()
.map(|neuron| neuron.controller.unwrap())
.collect();

nns_neurons
}
}

/// Installs the NNS canisters, ensuring that there is a whale neuron with `TEST_NEURON_1_ID`.
/// Requires PocketIC to have at least an NNS and an SNS subnet.
///
Expand All @@ -335,122 +538,23 @@ pub async fn install_nns_canisters(
custom_initial_registry_mutations: Option<Vec<RegistryAtomicMutateRequest>>,
neurons_fund_hotkeys: Vec<PrincipalId>,
) -> Vec<PrincipalId> {
let topology = pocket_ic.topology().await;

let sns_subnet_id = topology.get_sns().expect("No SNS subnet found");
let sns_subnet_id = PrincipalId::from(sns_subnet_id);
let sns_subnet_id = SubnetId::from(sns_subnet_id);
println!("sns_subnet_id = {:?}", sns_subnet_id);
let mut nns_installer = NnsInstaller::default();

let mut nns_init_payload_builder = NnsInitPayloadsBuilder::new();
nns_installer
.with_initial_balances(initial_balances)
.with_neurons_fund_hotkeys(neurons_fund_hotkeys);

if let Some(custom_initial_registry_mutations) = custom_initial_registry_mutations {
nns_init_payload_builder.with_initial_mutations(custom_initial_registry_mutations);
if with_mainnet_nns_canister_versions {
nns_installer.with_mainnet_nns_canister_versions();
} else {
nns_init_payload_builder.with_initial_invariant_compliant_mutations();
nns_installer.with_tip_nns_canister_versions();
}
let maturity_equivalent_icp_e8s = 1_500_000 * E8;
nns_init_payload_builder
.with_test_neurons_fund_neurons_with_hotkeys(
neurons_fund_hotkeys,
maturity_equivalent_icp_e8s,
)
.with_sns_dedicated_subnets(vec![sns_subnet_id])
.with_sns_wasm_access_controls(true);

for (test_user_icp_ledger_account, test_user_icp_ledger_initial_balance) in initial_balances {
nns_init_payload_builder.with_ledger_account(
test_user_icp_ledger_account,
test_user_icp_ledger_initial_balance,
);
if let Some(custom_initial_registry_mutations) = custom_initial_registry_mutations {
nns_installer.with_custom_initial_registry_mutations(custom_initial_registry_mutations);
}

let nns_init_payload = nns_init_payload_builder.build();

let (governance_wasm, ledger_wasm, root_wasm, lifeline_wasm, sns_wasm_wasm, registry_wasm) =
if with_mainnet_nns_canister_versions {
(
build_mainnet_governance_wasm(),
build_mainnet_ledger_wasm(),
build_mainnet_root_wasm(),
build_mainnet_lifeline_wasm(),
build_mainnet_sns_wasms_wasm(),
build_mainnet_registry_wasm(),
)
} else {
(
build_governance_wasm(),
build_ledger_wasm(),
build_root_wasm(),
build_lifeline_wasm(),
build_sns_wasms_wasm(),
build_registry_wasm(),
)
};

install_canister(
pocket_ic,
"ICP Ledger",
LEDGER_CANISTER_ID,
Encode!(&nns_init_payload.ledger).unwrap(),
ledger_wasm,
Some(ROOT_CANISTER_ID.get()),
)
.await;
install_canister(
pocket_ic,
"NNS Root",
ROOT_CANISTER_ID,
Encode!(&nns_init_payload.root).unwrap(),
root_wasm,
Some(LIFELINE_CANISTER_ID.get()),
)
.await;
install_canister(
pocket_ic,
"NNS Governance",
GOVERNANCE_CANISTER_ID,
nns_init_payload.governance.encode_to_vec(),
governance_wasm,
Some(ROOT_CANISTER_ID.get()),
)
.await;
install_canister(
pocket_ic,
"Lifeline",
LIFELINE_CANISTER_ID,
Encode!(&nns_init_payload.lifeline).unwrap(),
lifeline_wasm,
Some(ROOT_CANISTER_ID.get()),
)
.await;
install_canister(
pocket_ic,
"NNS SNS-W",
SNS_WASM_CANISTER_ID,
Encode!(&nns_init_payload.sns_wasms).unwrap(),
sns_wasm_wasm,
Some(ROOT_CANISTER_ID.get()),
)
.await;
install_canister(
pocket_ic,
"Registry",
REGISTRY_CANISTER_ID,
Encode!(&nns_init_payload.registry).unwrap(),
registry_wasm,
Some(ROOT_CANISTER_ID.get()),
)
.await;

let nns_neurons = nns_init_payload
.governance
.neurons
.values()
.map(|neuron| neuron.controller.unwrap())
.collect();

nns_neurons
nns_installer.install(pocket_ic).await
}

#[derive(Copy, Clone, Debug)]
Expand Down