stake-pool: Add test for removing all validators from huge pool (#3806)

Author: joncinque

stake-pool/program/tests/helpers/mod.rs

@@ -1706,6 +1706,7 @@ pub fn add_vote_account(program_test: &mut ProgramTest) -> Pubkey {
     vote_pubkey
 }
 
+#[allow(clippy::too_many_arguments)]
 pub fn add_validator_stake_account(
     program_test: &mut ProgramTest,
     stake_pool: &mut state::StakePool,
@@ -1714,6 +1715,7 @@ pub fn add_validator_stake_account(
     withdraw_authority: &Pubkey,
     voter_pubkey: &Pubkey,
     stake_amount: u64,
+    status: state::StakeStatus,
 ) {
     let meta = stake::state::Meta {
         rent_exempt_reserve: STAKE_ACCOUNT_RENT_EXEMPTION,
@@ -1756,7 +1758,7 @@ pub fn add_validator_stake_account(
     let active_stake_lamports = stake_amount - LAMPORTS_PER_SOL;
 
     validator_list.validators.push(state::ValidatorStakeInfo {
-        status: state::StakeStatus::Active,
+        status,
         vote_account_address: *voter_pubkey,
         active_stake_lamports,
         transient_stake_lamports: 0,

stake-pool/program/tests/huge_pool.rs

@@ -57,6 +57,7 @@ async fn setup(
             &stake_pool_accounts.withdraw_authority,
             vote_account_address,
             stake_amount,
+            StakeStatus::Active,
         );
     }
 
@@ -578,3 +579,62 @@ async fn withdraw() {
         .await;
     assert!(error.is_none(), "{:?}", error);
 }
+
+#[tokio::test]
+async fn cleanup_all() {
+    let mut program_test = program_test();
+    let mut vote_account_pubkeys = vec![];
+    let mut stake_pool_accounts = StakePoolAccounts::new();
+    let max_validators = HUGE_POOL_SIZE;
+    stake_pool_accounts.max_validators = max_validators;
+
+    let stake_pool_pubkey = stake_pool_accounts.stake_pool.pubkey();
+    let (mut stake_pool, mut validator_list) = stake_pool_accounts.state();
+
+    for _ in 0..max_validators {
+        vote_account_pubkeys.push(add_vote_account(&mut program_test));
+    }
+
+    for vote_account_address in vote_account_pubkeys.iter() {
+        add_validator_stake_account(
+            &mut program_test,
+            &mut stake_pool,
+            &mut validator_list,
+            &stake_pool_pubkey,
+            &stake_pool_accounts.withdraw_authority,
+            vote_account_address,
+            STAKE_AMOUNT,
+            StakeStatus::ReadyForRemoval,
+        );
+    }
+
+    add_stake_pool_account(
+        &mut program_test,
+        &stake_pool_accounts.stake_pool.pubkey(),
+        &stake_pool,
+    );
+    add_validator_list_account(
+        &mut program_test,
+        &stake_pool_accounts.validator_list.pubkey(),
+        &validator_list,
+        max_validators,
+    );
+    let mut context = program_test.start_with_context().await;
+
+    let transaction = Transaction::new_signed_with_payer(
+        &[instruction::cleanup_removed_validator_entries(
+            &id(),
+            &stake_pool_accounts.stake_pool.pubkey(),
+            &stake_pool_accounts.validator_list.pubkey(),
+        )],
+        Some(&context.payer.pubkey()),
+        &[&context.payer],
+        context.last_blockhash,
+    );
+    let error = context
+        .banks_client
+        .process_transaction(transaction)
+        .await
+        .err();
+    assert!(error.is_none());
+}