From 2572df745298e94ed515636b10afb511b60250f0 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <jeff.washington@solana.com>
Date: Sun, 11 Jun 2023 14:26:16 -0500
Subject: [PATCH]  stake tests modification: next_epoch_and_n_slots (#32052)

stake tests modification: next_epoch_and_n_slots
---
 runtime/tests/stake.rs | 29 +++++++++++++++++++++--------
 1 file changed, 21 insertions(+), 8 deletions(-)
 mode change 100644 => 100755 runtime/tests/stake.rs

diff --git a/runtime/tests/stake.rs b/runtime/tests/stake.rs
old mode 100644
new mode 100755
index 6a74d87f6f..629cb9998f
--- a/runtime/tests/stake.rs
+++ b/runtime/tests/stake.rs
@@ -29,14 +29,26 @@ use {
     std::sync::Arc,
 };
 
-fn next_epoch(bank: &Arc<Bank>) -> Arc<Bank> {
+/// get bank at next epoch + `n` slots
+fn next_epoch_and_n_slots(bank: &Arc<Bank>, mut n: usize) -> Arc<Bank> {
     bank.squash();
-
-    Arc::new(Bank::new_from_parent(
+    let mut bank = Arc::new(Bank::new_from_parent(
         bank,
         &Pubkey::default(),
         bank.get_slots_in_epoch(bank.epoch()) + bank.slot(),
-    ))
+    ));
+
+    while n > 0 {
+        bank.squash();
+        bank = Arc::new(Bank::new_from_parent(
+            &bank,
+            &Pubkey::default(),
+            1 + bank.slot(),
+        ));
+        n -= 1;
+    }
+
+    bank
 }
 
 fn fill_epoch_with_votes(
@@ -385,7 +397,7 @@ fn test_stake_account_lifetime() {
             break;
         }
         // Cycle thru banks until we're fully warmed up
-        bank = next_epoch(&bank);
+        bank = next_epoch_and_n_slots(&bank, 0);
     }
 
     // Reward redemption
@@ -409,7 +421,7 @@ fn test_stake_account_lifetime() {
     let pre_balance = bank.get_balance(&stake_pubkey);
 
     // next epoch bank should pay rewards
-    bank = next_epoch(&bank);
+    bank = next_epoch_and_n_slots(&bank, 0);
 
     // Test that balance increased, and that the balance got staked
     let staked = get_staked(&bank, &stake_pubkey);
@@ -477,7 +489,8 @@ fn test_stake_account_lifetime() {
         .send_and_confirm_message(&[&mint_keypair, &stake_keypair], message)
         .is_err());
 
-    let mut bank = next_epoch(&bank);
+    let mut bank = next_epoch_and_n_slots(&bank, 0);
+
     let bank_client = BankClient::new_shared(&bank);
 
     // assert we're still cooling down
@@ -522,7 +535,7 @@ fn test_stake_account_lifetime() {
         if get_staked(&bank, &split_stake_pubkey) == 0 {
             break;
         }
-        bank = next_epoch(&bank);
+        bank = next_epoch_and_n_slots(&bank, 0);
     }
     let bank_client = BankClient::new_shared(&bank);