Refactor the common test snapshot dirs creation code (#31223)

* Create fn create_snapshot_dirs_for_tests to refactor the common test snapshot dirs creation code

* Fix trailing space

* Fix missed set_accounts_delta_hash

runtime/src/snapshot_utils.rs

@@ -5117,25 +5117,25 @@ mod tests {
         assert!(matches!(ret, Err(SnapshotError::InvalidAppendVecPath(_))));
     }
 
-    #[test]
+    fn create_snapshot_dirs_for_tests(
-    fn test_get_highest_bank_snapshot() {
+        genesis_config: &GenesisConfig,
-        solana_logger::setup();
+        bank_snapshots_dir: impl AsRef<Path>,
+        num_total: usize,
+        num_posts: usize,
+    ) -> Bank {
+        let mut bank = Arc::new(Bank::new_for_tests(genesis_config));
 
-        let genesis_config = GenesisConfig::default();
-        let mut bank = Arc::new(Bank::new_for_tests(&genesis_config));
-
-        let bank_snapshots_dir = tempfile::TempDir::new().unwrap();
         let collecter_id = Pubkey::new_unique();
         let snapshot_version = SnapshotVersion::default();
 
-        for _ in 0..4 {
+        // loop to create the banks at slot 1 to num_total
+        for _ in 0..num_total {
             // prepare the bank
             bank = Arc::new(Bank::new_from_parent(&bank, &collecter_id, bank.slot() + 1));
             bank.fill_bank_with_ticks_for_tests();
             bank.squash();
             bank.force_flush_accounts_cache();
 
-            // generate the bank snapshot directory for slot+1
             let snapshot_storages = bank.get_snapshot_storages(None);
             let slot_deltas = bank.status_cache.read().unwrap().root_slot_deltas();
             add_bank_snapshot(
@@ -5146,8 +5146,34 @@ mod tests {
                 slot_deltas,
             )
             .unwrap();
+
+            if bank.slot() as usize > num_posts {
+                continue; // leave the snapshot dir at PRE stage
+            }
+
+            // Reserialize the snapshot dir to convert it from PRE to POST, because only the POST type can be used
+            // to construct a bank.
+            assert!(
+                crate::serde_snapshot::reserialize_bank_with_new_accounts_hash(
+                    &bank_snapshots_dir,
+                    bank.slot(),
+                    &AccountsHash(Hash::new_unique()),
+                    None
+                )
+            );
         }
 
+        Arc::try_unwrap(bank).unwrap()
+    }
+
+    #[test]
+    fn test_get_highest_bank_snapshot() {
+        solana_logger::setup();
+
+        let genesis_config = GenesisConfig::default();
+        let bank_snapshots_dir = tempfile::TempDir::new().unwrap();
+        let _bank = create_snapshot_dirs_for_tests(&genesis_config, &bank_snapshots_dir, 4, 0);
+
         let snapshot = get_highest_bank_snapshot(&bank_snapshots_dir).unwrap();
         assert_eq!(snapshot.slot, 4);
 
@@ -5223,31 +5249,10 @@ mod tests {
     #[test]
     fn test_clean_orphaned_account_snapshot_dirs() {
         solana_logger::setup();
+
         let genesis_config = GenesisConfig::default();
-        let mut bank = Arc::new(Bank::new_for_tests(&genesis_config));
         let bank_snapshots_dir = tempfile::TempDir::new().unwrap();
-        let collecter_id = Pubkey::new_unique();
+        let _bank = create_snapshot_dirs_for_tests(&genesis_config, &bank_snapshots_dir, 2, 0);
-        let snapshot_version = SnapshotVersion::default();
-
-        for _ in 0..2 {
-            // prepare the bank
-            bank = Arc::new(Bank::new_from_parent(&bank, &collecter_id, bank.slot() + 1));
-            bank.fill_bank_with_ticks_for_tests();
-            bank.squash();
-            bank.force_flush_accounts_cache();
-
-            // generate the bank snapshot directory for slot+1
-            let snapshot_storages = bank.get_snapshot_storages(None);
-            let slot_deltas = bank.status_cache.read().unwrap().root_slot_deltas();
-            add_bank_snapshot(
-                &bank_snapshots_dir,
-                &bank,
-                &snapshot_storages,
-                snapshot_version,
-                slot_deltas,
-            )
-            .unwrap();
-        }
 
         let snapshot_dir_slot_2 = bank_snapshots_dir.path().join("2");
         let accounts_link_dir_slot_2 = snapshot_dir_slot_2.join("accounts_hardlinks");
@@ -5456,32 +5461,10 @@ mod tests {
     #[test]
     fn test_bank_from_snapshot_dir() {
         solana_logger::setup();
+
         let genesis_config = GenesisConfig::default();
-        let mut bank = Arc::new(Bank::new_for_tests(&genesis_config));
-
         let bank_snapshots_dir = tempfile::TempDir::new().unwrap();
-        let collecter_id = Pubkey::new_unique();
+        let bank = create_snapshot_dirs_for_tests(&genesis_config, &bank_snapshots_dir, 3, 0);
-        let snapshot_version = SnapshotVersion::default();
-
-        for _ in 0..3 {
-            // prepare the bank
-            bank = Arc::new(Bank::new_from_parent(&bank, &collecter_id, bank.slot() + 1));
-            bank.fill_bank_with_ticks_for_tests();
-            bank.squash();
-            bank.force_flush_accounts_cache();
-
-            // generate the bank snapshot directory for slot+1
-            let snapshot_storages = bank.get_snapshot_storages(None);
-            let slot_deltas = bank.status_cache.read().unwrap().root_slot_deltas();
-            add_bank_snapshot(
-                &bank_snapshots_dir,
-                &bank,
-                &snapshot_storages,
-                snapshot_version,
-                slot_deltas,
-            )
-            .unwrap();
-        }
 
         let bank_snapshot = get_highest_bank_snapshot(&bank_snapshots_dir).unwrap();
         let account_paths = &bank.rc.accounts.accounts_db.paths;
@@ -5504,7 +5487,7 @@ mod tests {
         .unwrap();
 
         bank_constructed.wait_for_initial_accounts_hash_verification_completed_for_tests();
-        assert_eq!(bank_constructed, *bank);
+        assert_eq!(bank_constructed, bank);
 
         // Verify that the next_append_vec_id tracking is correct
         let mut max_id = 0;
@@ -5523,42 +5506,10 @@ mod tests {
     #[test]
     fn test_bank_from_latest_snapshot_dir() {
         solana_logger::setup();
+
         let genesis_config = GenesisConfig::default();
-        let mut bank = Arc::new(Bank::new_for_tests(&genesis_config));
-
         let bank_snapshots_dir = tempfile::TempDir::new().unwrap();
-        let collecter_id = Pubkey::new_unique();
+        let bank = create_snapshot_dirs_for_tests(&genesis_config, &bank_snapshots_dir, 3, 3);
-        let snapshot_version = SnapshotVersion::default();
-
-        for _ in 0..3 {
-            // prepare the bank
-            bank = Arc::new(Bank::new_from_parent(&bank, &collecter_id, bank.slot() + 1));
-            bank.fill_bank_with_ticks_for_tests();
-            bank.squash();
-            bank.force_flush_accounts_cache();
-
-            // generate the bank snapshot directory for slot+1
-            let snapshot_storages = bank.get_snapshot_storages(None);
-            let slot_deltas = bank.status_cache.read().unwrap().root_slot_deltas();
-            add_bank_snapshot(
-                &bank_snapshots_dir,
-                &bank,
-                &snapshot_storages,
-                snapshot_version,
-                slot_deltas,
-            )
-            .unwrap();
-            // Reserialize the snapshot dir to convert it from PRE to POST, because only the POST type can be used
-            // to construct a bank.
-            assert!(
-                crate::serde_snapshot::reserialize_bank_with_new_accounts_hash(
-                    &bank_snapshots_dir,
-                    bank.slot(),
-                    &AccountsHash(Hash::new_unique()),
-                    None
-                )
-            );
-        }
 
         let account_paths = &bank.rc.accounts.accounts_db.paths;
 
@@ -5580,7 +5531,7 @@ mod tests {
         .unwrap();
 
         assert_eq!(
-            deserialized_bank, *bank,
+            deserialized_bank, bank,
             "Ensure rebuilding bank from the highest snapshot dir results in the highest bank",
         );
     }