Add Bank::set_epoch_stakes_for_test() so that we can adjust epoch_sta…

…kes in tests easily. (#2535) * Add Bank::set_epoch_stakes_for_test() so that we can adjust epoch_stakes in tests easily. * Do not use set_epoch_stakes_for_test to remove epoch_stakes. * Fix a bad merge.
anza-xyz · Aug 16, 2024 · f9316f8 · f9316f8
1 parent 1871e2c
commit f9316f8
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 13 deletions.
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
@@ -1974,6 +1974,11 @@ impl Bank {
         }
     }
 
+    #[cfg(feature = "dev-context-only-utils")]
+    pub fn set_epoch_stakes_for_test(&mut self, epoch: Epoch, stakes: EpochStakes) {
+        self.epoch_stakes.insert(epoch, stakes);
+    }
+
     fn update_rent(&self) {
         self.update_sysvar_account(&sysvar::rent::id(), |account| {
             create_account(

diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs
@@ -110,7 +110,8 @@ use {
     solana_vote_program::{
         vote_instruction,
         vote_state::{
-            self, BlockTimestamp, Vote, VoteInit, VoteState, VoteStateVersions, MAX_LOCKOUT_HISTORY,
+            self, create_account_with_authorized, BlockTimestamp, Vote, VoteInit, VoteState,
+            VoteStateVersions, MAX_LOCKOUT_HISTORY,
         },
     },
     std::{
@@ -12996,7 +12997,7 @@ fn test_bank_epoch_stakes() {
         create_genesis_config_with_vote_accounts(1_000_000_000, &voting_keypairs, stakes.clone());
 
     let bank0 = Arc::new(Bank::new_for_tests(&genesis_config));
-    let bank1 = Bank::new_from_parent(
+    let mut bank1 = Bank::new_from_parent(
         bank0.clone(),
         &Pubkey::default(),
         bank0.get_slots_in_epoch(0) + 1,
@@ -13023,4 +13024,33 @@ fn test_bank_epoch_stakes() {
             Some(stakes[i])
         );
     }
+
+    let new_epoch_stakes = EpochStakes::new_for_tests(
+        voting_keypairs
+            .iter()
+            .map(|keypair| {
+                let node_id = keypair.node_keypair.pubkey();
+                let authorized_voter = keypair.vote_keypair.pubkey();
+                let vote_account = VoteAccount::try_from(create_account_with_authorized(
+                    &node_id,
+                    &authorized_voter,
+                    &node_id,
+                    0,
+                    100,
+                ))
+                .unwrap();
+                (authorized_voter, (100_u64, vote_account))
+            })
+            .collect::<HashMap<_, _>>(),
+        1,
+    );
+    bank1.set_epoch_stakes_for_test(1, new_epoch_stakes);
+    assert_eq!(bank1.epoch_total_stake(1), Some(100 * num_of_nodes));
+    assert_eq!(bank1.epoch_node_id_to_stake(1, &Pubkey::new_unique()), None);
+    for keypair in voting_keypairs.iter() {
+        assert_eq!(
+            bank1.epoch_node_id_to_stake(1, &keypair.node_keypair.pubkey()),
+            Some(100)
+        );
+    }
 }
diff --git a/runtime/src/epoch_stakes.rs b/runtime/src/epoch_stakes.rs
@@ -40,6 +40,21 @@ impl EpochStakes {
         }
     }
 
+    #[cfg(feature = "dev-context-only-utils")]
+    pub fn new_for_tests(
+        vote_accounts_hash_map: VoteAccountsHashMap,
+        leader_schedule_epoch: Epoch,
+    ) -> Self {
+        Self::new(
+            Arc::new(StakesEnum::Accounts(crate::stakes::Stakes::new_for_tests(
+                0,
+                solana_vote::vote_account::VoteAccounts::from(Arc::new(vote_accounts_hash_map)),
+                im::HashMap::default(),
+            ))),
+            leader_schedule_epoch,
+        )
+    }
+
     pub fn stakes(&self) -> &StakesEnum {
         &self.stakes
     }
@@ -228,10 +243,9 @@ pub(crate) mod tests {
             stake_account::StakeAccount,
             stakes::{Stakes, StakesCache},
         },
-        im::HashMap as ImHashMap,
         solana_sdk::{account::AccountSharedData, rent::Rent},
         solana_stake_program::stake_state::{self, Delegation, Stake},
-        solana_vote::vote_account::{VoteAccount, VoteAccounts},
+        solana_vote::vote_account::VoteAccount,
         solana_vote_program::vote_state::{self, create_account_with_authorized},
         std::iter,
     };
@@ -544,14 +558,7 @@ pub(crate) mod tests {
         let epoch_vote_accounts = new_epoch_vote_accounts(&vote_accounts_map, |node_id| {
             *node_id_to_stake_map.get(node_id).unwrap()
         });
-        let epoch_stakes = EpochStakes::new(
-            Arc::new(StakesEnum::Accounts(Stakes::new_for_tests(
-                0,
-                VoteAccounts::from(Arc::new(epoch_vote_accounts)),
-                ImHashMap::default(),
-            ))),
-            0,
-        );
+        let epoch_stakes = EpochStakes::new_for_tests(epoch_vote_accounts, 0);
 
         assert_eq!(epoch_stakes.total_stake(), 11000);
         for (node_id, stake) in node_id_to_stake_map.iter() {

diff --git a/runtime/src/stakes.rs b/runtime/src/stakes.rs
@@ -320,7 +320,7 @@ impl Stakes<StakeAccount> {
         })
     }
 
-    #[cfg(test)]
+    #[cfg(feature = "dev-context-only-utils")]
     pub fn new_for_tests(
         epoch: Epoch,
         vote_accounts: VoteAccounts,