test(sns): Add a few unit tests of refresh_cached_upgrade_steps and i…

…nitiate_upgrade_if_sns_behind_target_version (#2730)

What I wanted to cover was

1. Ensuring `initiate_upgrade_if_sns_behind_target_version` would be
properly blocked by other upgrade proposals and pending upgrades

3. Ensuring `refresh_cached_upgrade_steps` doesn't do anything crazy
when SNS-W gives a weird or erroneous response

rs/sns/governance/src/governance.rs

@@ -4827,6 +4827,7 @@ impl Governance {
             }
         }
     }
+
     /// Checks if an automatic upgrade is needed and initiates it.
     /// An automatic upgrade is needed if `target_version` is set to a future version on the upgrade path
     async fn initiate_upgrade_if_sns_behind_target_version(&mut self) {
@@ -6030,3 +6031,6 @@ mod assorted_governance_tests;
 
 #[cfg(test)]
 mod fail_stuck_upgrade_in_progress_tests;
+
+#[cfg(test)]
+mod advance_target_sns_version_tests;

rs/sns/governance/src/governance/advance_target_sns_version_tests.rs

@@ -0,0 +1,252 @@
+use super::assorted_governance_tests::{
+    basic_governance_proto, DoNothingLedger, TEST_GOVERNANCE_CANISTER_ID, TEST_ROOT_CANISTER_ID,
+};
+use super::*;
+use crate::{
+    pb::v1::{ProposalData, Tally, UpgradeSnsToNextVersion},
+    sns_upgrade::{ListUpgradeStep, ListUpgradeStepsRequest, ListUpgradeStepsResponse, SnsVersion},
+    types::test_helpers::NativeEnvironment,
+};
+use ic_nervous_system_common::cmc::FakeCmc;
+use ic_nns_constants::SNS_WASM_CANISTER_ID;
+use maplit::btreemap;
+use pretty_assertions::assert_eq;
+
+#[tokio::test]
+async fn test_initiate_upgrade_blocked_by_upgrade_proposal() {
+    // Step 1: Prepare the world.
+    let root_canister_id = *TEST_ROOT_CANISTER_ID;
+    let governance_canister_id = *TEST_GOVERNANCE_CANISTER_ID;
+
+    let mut env = NativeEnvironment::new(Some(governance_canister_id));
+
+    let current_version = SnsVersion {
+        root_wasm_hash: vec![1, 2, 3],
+        governance_wasm_hash: vec![2, 3, 4],
+        ledger_wasm_hash: vec![3, 4, 5],
+        swap_wasm_hash: vec![4, 5, 6],
+        archive_wasm_hash: vec![5, 6, 7],
+        index_wasm_hash: vec![6, 7, 8],
+    };
+
+    let target_version = {
+        let mut version = current_version.clone();
+        version.governance_wasm_hash = vec![9, 9, 9];
+        version
+    };
+
+    // Set up environment to return upgrade steps that would allow an upgrade
+    env.set_call_canister_response(
+        SNS_WASM_CANISTER_ID,
+        "list_upgrade_steps",
+        Encode!(&ListUpgradeStepsRequest {
+            starting_at: Some(current_version.clone()),
+            sns_governance_canister_id: Some(governance_canister_id.into()),
+            limit: 0,
+        })
+        .unwrap(),
+        Ok(Encode!(&ListUpgradeStepsResponse {
+            steps: vec![
+                ListUpgradeStep {
+                    version: Some(current_version.clone())
+                },
+                ListUpgradeStep {
+                    version: Some(target_version.clone())
+                },
+            ]
+        })
+        .unwrap()),
+    );
+
+    let proposal_id = 12;
+    let action = Action::UpgradeSnsToNextVersion(UpgradeSnsToNextVersion {});
+    let proposal = ProposalData {
+        action: (&action).into(),
+        id: Some(proposal_id.into()),
+        decided_timestamp_seconds: 1,
+        latest_tally: Some(Tally {
+            yes: 1,
+            no: 0,
+            total: 1,
+            timestamp_seconds: 1,
+        }),
+        ..Default::default()
+    };
+    let mut governance = Governance::new(
+        GovernanceProto {
+            root_canister_id: Some(root_canister_id.get()),
+            deployed_version: Some(current_version.clone().into()),
+            target_version: Some(target_version.clone().into()),
+            // Add an upgrade proposal that is adopted but not executed
+            proposals: btreemap! {
+                proposal_id => proposal
+            },
+            ..basic_governance_proto()
+        }
+        .try_into()
+        .unwrap(),
+        Box::new(env),
+        Box::new(DoNothingLedger {}),
+        Box::new(DoNothingLedger {}),
+        Box::new(FakeCmc::new()),
+    );
+
+    // Step 2: Run code under test.
+    assert_eq!(governance.proto.cached_upgrade_steps, None);
+    governance.temporarily_lock_refresh_cached_upgrade_steps();
+    governance.refresh_cached_upgrade_steps().await;
+    assert_eq!(
+        governance
+            .proto
+            .cached_upgrade_steps
+            .clone()
+            .unwrap()
+            .upgrade_steps
+            .unwrap()
+            .versions
+            .len(),
+        2
+    );
+
+    governance
+        .initiate_upgrade_if_sns_behind_target_version()
+        .await;
+
+    // Step 3: Inspect results.
+    // The pending_version should remain None since we had an upgrade proposal in progress
+    assert_eq!(governance.proto.pending_version, None);
+
+    // The deployed version should remain unchanged
+    assert_eq!(
+        governance.proto.deployed_version,
+        Some(Version::from(current_version))
+    );
+
+    // The target version should remain unchanged
+    assert_eq!(
+        governance.proto.target_version,
+        Some(Version::from(target_version))
+    );
+}
+
+#[tokio::test]
+async fn test_initiate_upgrade_blocked_by_pending_upgrade() {
+    // Step 1: Prepare the world.
+    let root_canister_id = *TEST_ROOT_CANISTER_ID;
+    let governance_canister_id = *TEST_GOVERNANCE_CANISTER_ID;
+
+    let mut env = NativeEnvironment::new(Some(governance_canister_id));
+
+    let current_version = SnsVersion {
+        root_wasm_hash: vec![1, 2, 3],
+        governance_wasm_hash: vec![2, 3, 4],
+        ledger_wasm_hash: vec![3, 4, 5],
+        swap_wasm_hash: vec![4, 5, 6],
+        archive_wasm_hash: vec![5, 6, 7],
+        index_wasm_hash: vec![6, 7, 8],
+    };
+
+    let target_version = {
+        let mut version = current_version.clone();
+        version.governance_wasm_hash = vec![9, 9, 9];
+        version
+    };
+
+    // Set up environment to return upgrade steps that would allow an upgrade
+    env.set_call_canister_response(
+        SNS_WASM_CANISTER_ID,
+        "list_upgrade_steps",
+        Encode!(&ListUpgradeStepsRequest {
+            starting_at: Some(current_version.clone()),
+            sns_governance_canister_id: Some(governance_canister_id.into()),
+            limit: 0,
+        })
+        .unwrap(),
+        Ok(Encode!(&ListUpgradeStepsResponse {
+            steps: vec![
+                ListUpgradeStep {
+                    version: Some(current_version.clone())
+                },
+                ListUpgradeStep {
+                    version: Some(target_version.clone())
+                },
+            ]
+        })
+        .unwrap()),
+    );
+
+    let pending_version = Version {
+        root_wasm_hash: vec![8, 8, 8],
+        governance_wasm_hash: vec![8, 8, 8],
+        ledger_wasm_hash: vec![8, 8, 8],
+        swap_wasm_hash: vec![8, 8, 8],
+        archive_wasm_hash: vec![8, 8, 8],
+        index_wasm_hash: vec![8, 8, 8],
+    };
+    let mut governance = Governance::new(
+        GovernanceProto {
+            root_canister_id: Some(root_canister_id.get()),
+            deployed_version: Some(Version::from(current_version.clone())),
+            target_version: Some(Version::from(target_version.clone())),
+            // There's already an upgrade pending
+            pending_version: Some(PendingVersion {
+                target_version: Some(pending_version.clone()),
+                mark_failed_at_seconds: 123,
+                checking_upgrade_lock: 0,
+                proposal_id: Some(42),
+            }),
+            ..basic_governance_proto()
+        }
+        .try_into()
+        .unwrap(),
+        Box::new(env),
+        Box::new(DoNothingLedger {}),
+        Box::new(DoNothingLedger {}),
+        Box::new(FakeCmc::new()),
+    );
+
+    // Step 2: Run code under test.
+    assert_eq!(governance.proto.cached_upgrade_steps, None);
+    governance.temporarily_lock_refresh_cached_upgrade_steps();
+    governance.refresh_cached_upgrade_steps().await;
+    assert_eq!(
+        governance
+            .proto
+            .cached_upgrade_steps
+            .clone()
+            .unwrap()
+            .upgrade_steps
+            .unwrap()
+            .versions
+            .len(),
+        2
+    );
+
+    governance
+        .initiate_upgrade_if_sns_behind_target_version()
+        .await;
+
+    // Step 3: Inspect results.
+    // The pending_version should remain unchanged since we had an upgrade in progress
+    assert_eq!(
+        governance.proto.pending_version,
+        Some(PendingVersion {
+            target_version: Some(pending_version),
+            mark_failed_at_seconds: 123,
+            checking_upgrade_lock: 0,
+            proposal_id: Some(42),
+        })
+    );
+
+    // The deployed version should remain unchanged
+    assert_eq!(
+        governance.proto.deployed_version,
+        Some(Version::from(current_version))
+    );
+
+    // The target version should remain unchanged
+    assert_eq!(
+        governance.proto.target_version,
+        Some(Version::from(target_version))
+    );
+}

rs/sns/governance/tests/governance.rs

@@ -10,6 +10,7 @@ use ic_nervous_system_common_test_keys::{
     TEST_NEURON_1_OWNER_PRINCIPAL, TEST_NEURON_2_OWNER_PRINCIPAL,
 };
 use ic_nervous_system_proto::pb::v1::{Percentage, Principals};
+use ic_sns_governance::pb::v1::governance::CachedUpgradeSteps;
 use ic_sns_governance::{
     governance::{
         MATURITY_DISBURSEMENT_DELAY_SECONDS, UPGRADE_STEPS_INTERVAL_REFRESH_BACKOFF_SECONDS,
@@ -3054,6 +3055,78 @@ async fn test_refresh_cached_upgrade_steps() {
     }
 }
 
+#[tokio::test]
+async fn test_refresh_cached_upgrade_steps_doesnt_panic_on_invalid_response() {
+    let mut canister_fixture = GovernanceCanisterFixtureBuilder::new().create();
+
+    // Set up the fixture state with a deployed version
+    canister_fixture.governance.proto.deployed_version = Some(Version::default());
+
+    // Mock SNS-W to return an invalid response (empty steps)
+    canister_fixture
+        .environment_fixture
+        .push_mocked_canister_reply(ListUpgradeStepsResponse { steps: vec![] });
+
+    // Initial state should be None
+    assert_eq!(canister_fixture.governance.proto.cached_upgrade_steps, None);
+
+    // Refresh should not panic on empty response
+    let now = canister_fixture.governance.env.now();
+    canister_fixture
+        .governance
+        .temporarily_lock_refresh_cached_upgrade_steps();
+    canister_fixture
+        .governance
+        .refresh_cached_upgrade_steps()
+        .await;
+    let expected_upgrade_steps = Some(CachedUpgradeSteps {
+        upgrade_steps: Some(Versions {
+            versions: Vec::new(),
+        }),
+        requested_timestamp_seconds: Some(now),
+        response_timestamp_seconds: Some(now),
+    });
+    assert_eq!(
+        canister_fixture.governance.proto.cached_upgrade_steps,
+        expected_upgrade_steps
+    );
+}
+
+#[tokio::test]
+async fn test_refresh_cached_upgrade_steps_handles_sns_w_error() {
+    let mut canister_fixture = GovernanceCanisterFixtureBuilder::new().create();
+
+    // Set up the fixture state with a deployed version
+    canister_fixture.governance.proto.deployed_version = Some(Version::default());
+
+    // Mock SNS-W to return an error
+    canister_fixture
+        .environment_fixture
+        .push_mocked_canister_panic("SNS-W error response");
+
+    let now = canister_fixture.governance.env.now();
+    let expected_upgrade_steps = Some(CachedUpgradeSteps {
+        upgrade_steps: Some(Versions {
+            versions: Vec::new(),
+        }),
+        requested_timestamp_seconds: Some(now),
+        response_timestamp_seconds: Some(now),
+    });
+    canister_fixture.governance.proto.cached_upgrade_steps = expected_upgrade_steps.clone();
+
+    // Refresh should not panic on error response
+    canister_fixture
+        .governance
+        .refresh_cached_upgrade_steps()
+        .await;
+
+    // State should remain None after error
+    assert_eq!(
+        canister_fixture.governance.proto.cached_upgrade_steps,
+        expected_upgrade_steps
+    );
+}
+
 #[tokio::test]
 async fn test_process_proposals_tallies_votes_for_proposals_where_voting_is_possible() {
     let mut canister_fixture = GovernanceCanisterFixtureBuilder::new()