refactor(mainnet-nns): fix regressions on testnet with mainnet state + refactor NNS recovery testnet

rs/tests/driver/src/driver/test_env_api.rs

@@ -1881,11 +1881,20 @@ pub trait HasPublicApiUrl: HasTestEnv + Send + Sync {
     }
 
     async fn await_status_is_healthy_async(&self) -> Result<()> {
+        self.await_status_is_healthy_with_retries_async(READY_WAIT_TIMEOUT, RETRY_BACKOFF)
+            .await
+    }
+
+    async fn await_status_is_healthy_with_retries_async(
+        &self,
+        timeout: Duration,
+        backoff: Duration,
+    ) -> Result<()> {
         retry_with_msg_async!(
             &format!("await_status_is_healthy of {}", self.get_public_url()),
             &self.test_env().logger(),
-            READY_WAIT_TIMEOUT,
-            RETRY_BACKOFF,
+            timeout,
+            backoff,
             || async {
                 self.status_is_healthy_async()
                     .await

rs/tests/nested/nns_recovery/common.rs

@@ -191,6 +191,16 @@ pub fn setup(env: TestEnv, cfg: SetupConfig) {
     )
     .setup_and_start(&env)
     .unwrap();
+
+    nested::registration(env.clone());
+    replace_nns_with_unassigned_nodes(&env);
+
+    let SshKeys {
+        ssh_priv_key_path: _,
+        auth: backup_auth,
+        ssh_pub_key: ssh_backup_pub_key,
+    } = get_ssh_keys_for_user(&env, BACKUP_USERNAME);
+    grant_backup_access_to_all_nns_nodes(&env, &backup_auth, &ssh_backup_pub_key);
 }
 
 pub fn test(env: TestEnv, cfg: TestConfig) {
@@ -205,14 +215,10 @@ pub fn test(env: TestEnv, cfg: TestConfig) {
     } = get_ssh_keys_for_user(&env, SSH_USERNAME);
     let SshKeys {
         ssh_priv_key_path: ssh_backup_priv_key_path,
-        auth: backup_auth,
-        ssh_pub_key: ssh_backup_pub_key,
+        auth: _,
+        ssh_pub_key: _,
     } = get_ssh_keys_for_user(&env, BACKUP_USERNAME);
 
-    nested::registration(env.clone());
-    replace_nns_with_unassigned_nodes(&env);
-    grant_backup_access_to_all_nns_nodes(&env, &backup_auth, &ssh_backup_pub_key);
-
     let current_version = get_guestos_img_version();
     info!(logger, "Current GuestOS version: {:?}", current_version);
 

rs/tests/testnets/BUILD.bazel

@@ -355,7 +355,6 @@ system_test_nns(
 system_test_nns(
     name = "nns_recovery",
     enable_metrics = True,
-    guestos_update = "test",
     setupos = True,
     tags = [
         "dynamic_testnet",
@@ -367,15 +366,14 @@ system_test_nns(
         "//rs/limits",
         "//rs/tests/consensus/subnet_recovery:lib",
         "//rs/tests/driver:ic-system-test-driver",
-        "//rs/tests/nested",
         "//rs/tests/nested/nns_recovery:common",
         "@crate_index//:anyhow",
         "@crate_index//:slog",
     ],
 )
 
 system_test_nns(
-    name = "mainnet_nns",
+    name = "mainnet_nns_state",
     enable_metrics = True,
     env = MAINNET_ENV,
     env_inherit = ["SSH_AUTH_SOCK"],

rs/tests/testnets/Cargo.toml

@@ -101,5 +101,5 @@ name = "nns_recovery"
 path = "nns_recovery.rs"
 
 [[bin]]
-name = "mainnet_nns"
-path = "mainnet_nns.rs"
+name = "mainnet_nns_state"
+path = "mainnet_nns_state.rs"

rs/tests/testnets/mainnet_nns/BUILD.bazel

@@ -15,7 +15,6 @@ rust_library(
         "//rs/canister_client",
         "//rs/canister_client/sender",
         "//rs/crypto/utils/threshold_sig_der",
-        "//rs/limits",
         "//rs/nervous_system/common",
         "//rs/nns/common",
         "//rs/nns/constants",

rs/tests/testnets/mainnet_nns/Cargo.toml

@@ -13,7 +13,6 @@ ic-base-types = { path = "../../../types/base_types" }
 ic-canister-client = { path = "../../../canister_client" }
 ic-canister-client-sender = { path = "../../../canister_client/sender" }
 ic-crypto-utils-threshold-sig-der = { path = "../../../crypto/utils/threshold_sig_der" }
-ic-limits = { path = "../../../limits" }
 ic-nervous-system-common = { path = "../../../nervous_system/common" }
 ic-nns-common = { path = "../../../nns/common" }
 ic-nns-constants = { path = "../../../nns/constants" }

rs/tests/testnets/mainnet_nns/src/lib.rs

@@ -2,7 +2,6 @@ use anyhow::Result;
 use ic_base_types::PrincipalId;
 use ic_consensus_system_test_utils::rw_message::install_nns_and_check_progress;
 use ic_crypto_utils_threshold_sig_der::public_key_der_to_pem;
-use ic_limits::DKG_INTERVAL_HEIGHT;
 use ic_nervous_system_common::E8;
 use ic_nns_common::types::NeuronId;
 use ic_registry_subnet_type::SubnetType;
@@ -79,7 +78,7 @@ pub mod proposals;
 /// except the root subnet (tdb26), which will contain only the single-node NNS subnet.
 /// Proposals can be made (and will instantly execute) using the relevant functions in
 /// `crate::proposals`.
-pub fn setup(env: TestEnv) {
+pub fn setup(env: TestEnv, dkg_interval: Option<u64>) {
     // Since we're creating the IC concurrently with fetching the state we use a channel to tell the
     // thread fetching the state when the IC is ready such that it can scp the ic.json5 config file
     // from the NNS node once it's online, used by ic-replay.
@@ -94,7 +93,7 @@ pub fn setup(env: TestEnv) {
     // Recover the NNS concurrently:
     let env_clone = env.clone();
     let recover_nns_thread = std::thread::spawn(move || {
-        setup_recovered_nns(env_clone, rx_finished_ic_setup, rx_aux_node)
+        setup_recovered_nns(env_clone, dkg_interval, rx_finished_ic_setup, rx_aux_node)
     });
 
     // Setup and start the aux UVM concurrently:
@@ -166,6 +165,7 @@ pub fn setup(env: TestEnv) {
 
 fn setup_recovered_nns(
     env: TestEnv,
+    dkg_interval: Option<u64>,
     rx_finished_ic_setup: Receiver<()>,
     rx_aux_node: Receiver<DeployedUniversalVm>,
 ) -> NeuronId {
@@ -211,46 +211,45 @@ fn setup_recovered_nns(
         neuron_id,
     );
 
-    let dkg_interval = std::env::var("DKG_INTERVAL")
-        .ok()
-        .and_then(|s| s.parse::<u64>().ok())
-        .unwrap_or(DKG_INTERVAL_HEIGHT);
-    let subnet_config = UpdateSubnetPayload {
-        subnet_id: SubnetId::from(PrincipalId::from_str(ORIGINAL_NNS_ID).unwrap()),
-        max_ingress_bytes_per_message: None,
-        max_ingress_messages_per_block: None,
-        max_ingress_bytes_per_block: None,
-        max_block_payload_size: None,
-        unit_delay_millis: None,
-        initial_notary_delay_millis: None,
-        dkg_interval_length: Some(dkg_interval),
-        dkg_dealings_per_block: None,
-        start_as_nns: None,
-        subnet_type: None,
-        is_halted: None,
-        halt_at_cup_height: None,
-        features: None,
-        resource_limits: None,
-        chain_key_config: None,
-        chain_key_signing_enable: None,
-        chain_key_signing_disable: None,
-        max_number_of_canisters: None,
-        ssh_readonly_access: None,
-        ssh_backup_access: None,
-        max_artifact_streams_per_peer: None,
-        max_chunk_wait_ms: None,
-        max_duplicity: None,
-        max_chunk_size: None,
-        receive_check_cache_size: None,
-        pfn_evaluation_period_ms: None,
-        registry_poll_period_ms: None,
-        retransmission_request_ms: None,
-        set_gossip_config_to_default: false,
-    };
-    block_on(ProposalWithMainnetState::update_subnet_record(
-        recovered_nns_node.get_public_url(),
-        subnet_config,
-    ));
+    if let Some(dkg_interval) = dkg_interval {
+        info!(env.logger(), "Overriding DKG interval to {dkg_interval}");
+        let subnet_config = UpdateSubnetPayload {
+            subnet_id: SubnetId::from(PrincipalId::from_str(ORIGINAL_NNS_ID).unwrap()),
+            max_ingress_bytes_per_message: None,
+            max_ingress_messages_per_block: None,
+            max_ingress_bytes_per_block: None,
+            max_block_payload_size: None,
+            unit_delay_millis: None,
+            initial_notary_delay_millis: None,
+            dkg_interval_length: Some(dkg_interval),
+            dkg_dealings_per_block: None,
+            start_as_nns: None,
+            subnet_type: None,
+            is_halted: None,
+            halt_at_cup_height: None,
+            features: None,
+            resource_limits: None,
+            chain_key_config: None,
+            chain_key_signing_enable: None,
+            chain_key_signing_disable: None,
+            max_number_of_canisters: None,
+            ssh_readonly_access: None,
+            ssh_backup_access: None,
+            max_artifact_streams_per_peer: None,
+            max_chunk_wait_ms: None,
+            max_duplicity: None,
+            max_chunk_size: None,
+            receive_check_cache_size: None,
+            pfn_evaluation_period_ms: None,
+            registry_poll_period_ms: None,
+            retransmission_request_ms: None,
+            set_gossip_config_to_default: false,
+        };
+        block_on(ProposalWithMainnetState::update_subnet_record(
+            recovered_nns_node.get_public_url(),
+            subnet_config,
+        ));
+    }
 
     let recovered_nns_pub_key = fetch_recovered_nns_public_key_pem(&recovered_nns_node);
 
@@ -479,8 +478,8 @@ fn recover_nns_subnet(
     let nns_ip = nns_node.get_ip_addr();
     let upload_ip = recovered_nns_node.get_ip_addr();
 
-    let recovery_dir = tempdir().unwrap().path().to_path_buf();
-    let mut cmd = Command::new(get_dependency_path_from_env("IC_RECOVERY_PATH"));
+    let recovery_dir = env.base_path();
+    let mut cmd = Command::new(env.get_path(PATH_IC_RECOVERY));
     cmd.arg("--skip-prompts")
         .arg("--dir")
         .arg(recovery_dir)
@@ -500,8 +499,10 @@ fn recover_nns_subnet(
         .arg(aux_ip.to_string())
         .arg("--aux-user")
         .arg(SSH_USERNAME)
+        // --validate-nns-url will not actually be used because the backup pod tarball also
+        // contains the local store, which initializes ic-recovery
         .arg("--validate-nns-url")
-        .arg(nns_url.to_string())
+        .arg("https://icp0.io")
         .arg("--upload-method")
         .arg(upload_ip.to_string())
         .arg("--parent-nns-host-ip")
@@ -580,7 +581,7 @@ fn patch_api_bn(env: &TestEnv, recovered_nns_node: &IcNodeSnapshot, api_bn: &IcN
     )
     .expect("Could not patch config NNS URLs of API BN");
 
-    // Path config NNS public key to the recovered NNS public key
+    // Patch config NNS public key to the recovered NNS public key
     patch_config_nns_public_key(
         &logger,
         api_bn,
@@ -604,8 +605,8 @@ fn patch_api_bn(env: &TestEnv, recovered_nns_node: &IcNodeSnapshot, api_bn: &IcN
         )
         .expect("Could not restart ic-replica on API BN");
 
-    api_bn
-        .await_status_is_healthy()
+    // Replicating mainnet registry can take a bit of time
+    block_on(api_bn.await_status_is_healthy_with_retries_async(secs(15 * 60), secs(30)))
         .expect("API BN did not become healthy after patching");
 }
 
@@ -706,7 +707,7 @@ fn write_sh_lib(env: &TestEnv, neuron_id: NeuronId, http_gateway: &Url) {
     let logger: slog::Logger = env.logger();
     let set_testnet_env_vars_sh_path = env.get_path(PATH_SET_TESTNET_ENV_VARS_SH);
     let set_testnet_env_vars_sh_str = set_testnet_env_vars_sh_path.display();
-    let ic_admin = fs::canonicalize(get_dependency_path_from_env("IC_ADMIN_PATH")).unwrap();
+    let ic_admin = fs::canonicalize(env.get_path("recovery/binaries/ic-admin")).unwrap();
     let pem = env.get_path("neuron_secret_key.pem");
     let mut pem_file = File::create(&pem).unwrap();
     pem_file
@@ -791,6 +792,6 @@ fn remove_large_files(env: &TestEnv) {
     let mut rm = Command::new("rm");
     rm.arg("-rf")
         .arg(env.get_path(PATH_STATE_TARBALL))
-        .arg(env.get_path("recovery"));
+        .arg(env.get_path(PATH_RECOVERY_WORKING_DIR));
     rm.output().expect("Failed to remove large files");
 }

rs/tests/testnets/mainnet_nns.rs → rs/tests/testnets/mainnet_nns_state.rs

@@ -24,7 +24,7 @@
 // ```
 //
 // Additional configuration:
-//  - --test_env=DKG_INTERVAL=<dkg_interval> (default: 499)
+//  - --test_env=DKG_INTERVAL=<dkg_interval>
 //  - --test_env=NNS_STATE_ON_BACKUP_POD=<path_on_remote_backup_pod> (default: dev@zh1-pyr07.zh1.dfinity.network:/home/dev/nns_state.tar.zst)
 //
 // To get access to P8s and Grafana look for the following lines in the ict console output:
@@ -41,9 +41,13 @@ use ic_testnet_mainnet_nns::setup as setup_mainnet_nns;
 use std::time::Duration;
 
 fn main() -> Result<()> {
+    let dkg_interval = std::env::var("DKG_INTERVAL")
+        .ok()
+        .map(|s| s.parse::<u64>().expect("DKG_INTERVAL must be a valid u64"));
+
     SystemTestGroup::new()
         .with_timeout_per_test(Duration::from_secs(90 * 60))
-        .with_setup(setup_mainnet_nns)
+        .with_setup(move |env| setup_mainnet_nns(env, dkg_interval))
         .execute_from_args()?;
     Ok(())
 }

rs/tests/testnets/nns_recovery.rs

@@ -31,14 +31,10 @@
 
 use anyhow::Result;
 use ic_consensus_system_test_subnet_recovery::utils::{
-    BACKUP_USERNAME, SshKeys, break_nodes, get_ssh_keys_for_user,
-    node_with_highest_certification_share_height,
+    break_nodes, node_with_highest_certification_share_height,
 };
 use ic_limits::DKG_INTERVAL_HEIGHT;
-use ic_nested_nns_recovery_common::{
-    NNS_RECOVERY_VM_RESOURCE_OVERRIDES, SetupConfig, grant_backup_access_to_all_nns_nodes,
-    replace_nns_with_unassigned_nodes,
-};
+use ic_nested_nns_recovery_common::{NNS_RECOVERY_VM_RESOURCE_OVERRIDES, SetupConfig};
 use ic_system_test_driver::driver::nested::HasNestedVms;
 use ic_system_test_driver::driver::test_env::{TestEnv, TestEnvAttribute};
 use ic_system_test_driver::driver::test_env_api::*;
@@ -55,7 +51,7 @@ fn setup(env: TestEnv) {
 
     let dkg_interval = std::env::var("DKG_INTERVAL")
         .ok()
-        .and_then(|s| s.parse::<u64>().ok())
+        .map(|s| s.parse::<u64>().expect("DKG_INTERVAL must be a valid u64"))
         .unwrap_or(DKG_INTERVAL_HEIGHT);
 
     ic_nested_nns_recovery_common::setup(
@@ -91,27 +87,6 @@ fn log_instructions(env: TestEnv) {
         );
     }
 
-    let SshKeys {
-        ssh_priv_key_path: _,
-        auth: backup_auth,
-        ssh_pub_key: ssh_backup_pub_key,
-    } = get_ssh_keys_for_user(&env, BACKUP_USERNAME);
-
-    nested::registration(env.clone());
-    replace_nns_with_unassigned_nodes(&env);
-    grant_backup_access_to_all_nns_nodes(&env, &backup_auth, &ssh_backup_pub_key);
-
-    let upgrade_version = get_guestos_update_img_version();
-    let upgrade_image_url = get_guestos_update_img_url();
-    let upgrade_image_hash = get_guestos_update_img_sha256();
-    info!(
-        logger,
-        r#"Working GuestOS version:
-    --upgrade-version {upgrade_version}
-    --upgrade-image-url {upgrade_image_url}
-    --upgrade-image-hash {upgrade_image_hash}"#
-    );
-
     info!(logger, "Host <-> IPs mapping:");
     for vm in env.get_all_nested_vms().unwrap() {
         let vm_name = vm.vm_name();