From e33c60d38a807e6c6c04160cc9d078098930bcf7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mathias=20Bj=C3=B6rkqvist?= <mathias.bjorkqvist@dfinity.org>
Date: Thu, 6 Feb 2025 11:01:32 +0100
Subject: [PATCH] Fix ICP golden state test

---
 rs/ledger_suite/icp/tests/golden_nns_state.rs | 25 +++++++++++++++----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/rs/ledger_suite/icp/tests/golden_nns_state.rs b/rs/ledger_suite/icp/tests/golden_nns_state.rs
index 5249d99834e..c6c0ab79285 100644
--- a/rs/ledger_suite/icp/tests/golden_nns_state.rs
+++ b/rs/ledger_suite/icp/tests/golden_nns_state.rs
@@ -253,7 +253,10 @@ fn should_create_state_machine_with_golden_nns_state() {
     setup.perform_upgrade_downgrade_testing(true);
 
     // Downgrade all the canisters to the mainnet version
-    setup.downgrade_to_mainnet();
+    // For breaking changes, e.g., if mainnet is running a version with balances and allowances in
+    // stable structures, but master also has blocks in stable structures, `ledger_is_downgradable`
+    // should be set to `false`, otherwise `true`.
+    setup.downgrade_to_mainnet(true);
 
     // Verify ledger balance and allowance state
     // As before, the allowance check needs to be skipped for the mainnet version of the ledger.
@@ -315,14 +318,20 @@ impl Setup {
         self.upgrade_archive_canisters(&self.master_wasms.archive);
     }
 
-    pub fn downgrade_to_mainnet(&self) {
+    pub fn downgrade_to_mainnet(&self, ledger_is_downgradable: bool) {
         println!("Downgrading to mainnet version");
         self.upgrade_index(&self.mainnet_wasms.index);
-        match self.upgrade_ledger(&self.mainnet_wasms.ledger) {
-            Ok(_) => {
+        match (
+            ledger_is_downgradable,
+            self.upgrade_ledger(&self.mainnet_wasms.ledger),
+        ) {
+            (false, Ok(_)) => {
                 panic!("should fail to downgrade ledger to mainnet version");
             }
-            Err(err) => {
+            (true, Ok(_)) => {
+                // The ledger was downgraded successfully, which is what was expected to happen
+            }
+            (false, Err(err)) => {
                 // The ledger will still be running the master version, while the other canisters
                 // will be downgraded to the mainnet version. This is not an ideal situation, nor
                 // is it expected to happen in practice, but for the moment let's run the test to
@@ -332,6 +341,12 @@ impl Setup {
                     "Trying to downgrade from incompatible version",
                 );
             }
+            (true, Err(err)) => {
+                panic!(
+                    "should successfully downgrade ledger to mainnet version: {}",
+                    err
+                );
+            }
         }
         self.check_ledger_metrics(ExpectMigration::No);
         self.upgrade_archive_canisters(&self.mainnet_wasms.archive);