fix test

anza-xyz · Jul 25, 2024 · 567973f · 567973f
1 parent 6765807
commit 567973f
Showing 1 changed file with 14 additions and 20 deletions.
diff --git a/runtime/src/bank/serde_snapshot.rs b/runtime/src/bank/serde_snapshot.rs
@@ -38,7 +38,6 @@ mod tests {
             pubkey::Pubkey, stake::state::Stake,
         },
         std::{
-            collections::HashMap,
             io::{BufReader, BufWriter, Cursor},
             ops::RangeFull,
             path::Path,
@@ -304,6 +303,19 @@ mod tests {
 
         // Set extra fields
         bank.fee_rate_governor.lamports_per_signature = 7000;
+        // Note that epoch_stakes already has two epoch stakes entries for epochs 0 and 1
+        // which will also be serialized to the versioned epoch stakes extra field. Those
+        // entries are of type Stakes<StakeAccount> so add a new entry for Stakes<Stake>.
+        bank.epoch_stakes.insert(
+            42,
+            EpochStakes::from(VersionedEpochStakes::Current {
+                stakes: Stakes::<Stake>::default(),
+                total_stake: 42,
+                node_id_to_vote_accounts: Arc::<NodeIdToVoteAccounts>::default(),
+                epoch_authorized_voters: Arc::<EpochAuthorizedVoters>::default(),
+            }),
+        );
+        assert_eq!(bank.epoch_stakes.len(), 3);
 
         // Serialize
         let snapshot_storages = bank.get_snapshot_storages(None);
@@ -317,18 +329,6 @@ mod tests {
         )
         .unwrap();
 
-        let mut new_epoch_stakes: HashMap<u64, VersionedEpochStakes> = HashMap::new();
-        new_epoch_stakes.insert(
-            42,
-            VersionedEpochStakes::Current {
-                stakes: Stakes::<Stake>::default(),
-                total_stake: 42,
-                node_id_to_vote_accounts: Arc::<NodeIdToVoteAccounts>::default(),
-                epoch_authorized_voters: Arc::<EpochAuthorizedVoters>::default(),
-            },
-        );
-        bincode::serialize_into(&mut writer, &new_epoch_stakes).unwrap();
-
         // Deserialize
         let rdr = Cursor::new(&buf[..]);
         let mut reader = std::io::BufReader::new(&buf[rdr.position() as usize..]);
@@ -362,13 +362,7 @@ mod tests {
         )
         .unwrap();
 
-        assert_eq!(
-            dbank.epoch_stakes(42),
-            Some(&EpochStakes::from(
-                new_epoch_stakes.get(&42).unwrap().clone()
-            ))
-        );
-
+        assert_eq!(bank.epoch_stakes, dbank.epoch_stakes);
         assert_eq!(
             bank.fee_rate_governor.lamports_per_signature,
             dbank.fee_rate_governor.lamports_per_signature