Compare ledger state before and after upgrade

dfinity · Sep 24, 2024 · 18a5b26 · 18a5b26
1 parent f5de7c7
commit 18a5b26
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 13 deletions.
diff --git a/rs/rosetta-api/icrc1/ledger/sm-tests/src/in_memory_ledger.rs b/rs/rosetta-api/icrc1/ledger/sm-tests/src/in_memory_ledger.rs
@@ -70,9 +70,11 @@ pub trait InMemoryLedgerState {
     fn validate_invariants(&self);
 }
 
+#[derive(Eq, PartialEq, Debug)]
 pub struct InMemoryLedger<K, AccountId, Tokens>
 where
-    K: Ord,
+    K: Ord + Hash,
+    AccountId: Hash + Eq,
 {
     balances: HashMap<AccountId, Tokens>,
     allowances: HashMap<K, Allowance<Tokens>>,
@@ -83,9 +85,9 @@ where
 
 impl<K, AccountId, Tokens> InMemoryLedgerState for InMemoryLedger<K, AccountId, Tokens>
 where
-    K: Ord + for<'a> From<(&'a AccountId, &'a AccountId)> + Clone + Hash,
+    K: Eq + PartialEq + Ord + for<'a> From<(&'a AccountId, &'a AccountId)> + Clone + Hash,
     K: Into<(AccountId, AccountId)>,
-    AccountId: PartialEq + Ord + Clone + Hash + std::fmt::Debug,
+    AccountId: Eq + PartialEq + Ord + Clone + Hash + std::fmt::Debug,
     Tokens: TokensType + Default,
 {
     type AccountId = AccountId;
@@ -524,7 +526,7 @@ impl InMemoryLedger<ApprovalKey, Account, Tokens> {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Eq, PartialEq, Clone, Debug)]
 pub struct BurnsWithoutSpender<AccountId> {
     pub minter: AccountId,
     pub burn_indexes: Vec<usize>,

diff --git a/rs/rosetta-api/icrc1/tests/golden_state_upgrade_downgrade.rs b/rs/rosetta-api/icrc1/tests/golden_state_upgrade_downgrade.rs
@@ -3,7 +3,7 @@ use candid::{Encode, Nat};
 use canister_test::Wasm;
 use ic_base_types::{CanisterId, PrincipalId};
 use ic_icrc1_ledger_sm_tests::in_memory_ledger::{
-    verify_ledger_state, ApprovalKey, BurnsWithoutSpender, InMemoryLedger, InMemoryLedgerState,
+    ApprovalKey, BurnsWithoutSpender, InMemoryLedger, InMemoryLedgerState,
 };
 use ic_icrc1_ledger_sm_tests::{
     get_all_ledger_and_archive_blocks, send_approval, send_transfer, send_transfer_from,
@@ -262,11 +262,31 @@ fn should_upgrade_icrc_sns_canisters_with_golden_state() {
     );
 }
 
+fn build_in_memory_ledger_and_compare_to_previous(
+    state_machine: &StateMachine,
+    canister_id: CanisterId,
+    burns_without_spender: Option<BurnsWithoutSpender<Account>>,
+    previous_ledger_state: Option<InMemoryLedger<ApprovalKey, Account, Tokens>>,
+) -> InMemoryLedger<ApprovalKey, Account, Tokens> {
+    let blocks = get_all_ledger_and_archive_blocks(state_machine, canister_id);
+    let mut in_memory_ledger_state =
+        ic_icrc1_ledger_sm_tests::in_memory_ledger::InMemoryLedger::new(burns_without_spender);
+    in_memory_ledger_state.ingest_icrc1_ledger_blocks(&blocks);
+    in_memory_ledger_state.verify_balances_and_allowances(state_machine, canister_id);
+    if let Some(previous_ledger_state) = previous_ledger_state {
+        assert_eq!(previous_ledger_state, in_memory_ledger_state);
+    }
+    in_memory_ledger_state
+}
+
 fn generate_transactions(
     state_machine: &StateMachine,
     canister_id: CanisterId,
-    in_memory_ledger: &mut InMemoryLedger<ApprovalKey, Account, Tokens>,
+    in_memory_ledger: &mut Option<InMemoryLedger<ApprovalKey, Account, Tokens>>,
 ) {
+    let Some(in_memory_ledger) = in_memory_ledger else {
+        return;
+    };
     let start = Instant::now();
     let minter_account = ic_icrc1_ledger_sm_tests::minting_account(state_machine, canister_id)
         .unwrap_or_else(|| panic!("minter account should be set for {:?}", canister_id));
@@ -520,8 +540,15 @@ fn perform_upgrade_downgrade_testing(
         );
         let canister_id =
             CanisterId::unchecked_from_principal(PrincipalId::from_str(canister_id_str).unwrap());
+        let mut previous_ledger_state = None;
         if extended_testing {
-            verify_ledger_state(state_machine, canister_id, burns_without_spender.clone());
+            previous_ledger_state = Some(build_in_memory_ledger_and_compare_to_previous(
+                state_machine,
+                canister_id,
+                burns_without_spender.clone(),
+                None,
+            ));
+            generate_transactions(state_machine, canister_id, &mut previous_ledger_state);
         }
         upgrade_canister(
             state_machine,
@@ -535,11 +562,13 @@ fn perform_upgrade_downgrade_testing(
             bump_gzip_timestamp(&master_canister_wasm),
         );
         if extended_testing {
-            let blocks = get_all_ledger_and_archive_blocks(state_machine, canister_id);
-            let mut expected_ledger_state =
-                ic_icrc1_ledger_sm_tests::in_memory_ledger::InMemoryLedger::new(None);
-            expected_ledger_state.ingest_icrc1_ledger_blocks(&blocks);
-            generate_transactions(state_machine, canister_id, &mut expected_ledger_state);
+            previous_ledger_state = Some(build_in_memory_ledger_and_compare_to_previous(
+                state_machine,
+                canister_id,
+                burns_without_spender.clone(),
+                previous_ledger_state,
+            ));
+            generate_transactions(state_machine, canister_id, &mut previous_ledger_state);
         }
         // Downgrade back to the mainnet ledger version
         upgrade_canister(
@@ -548,7 +577,13 @@ fn perform_upgrade_downgrade_testing(
             mainnet_canister_wasm.clone(),
         );
         if extended_testing {
-            verify_ledger_state(state_machine, canister_id, burns_without_spender);
+            previous_ledger_state = Some(build_in_memory_ledger_and_compare_to_previous(
+                state_machine,
+                canister_id,
+                burns_without_spender,
+                previous_ledger_state,
+            ));
+            generate_transactions(state_machine, canister_id, &mut previous_ledger_state);
         }
     }
 }