stake-pool: Split up more tests to help CI, avoid warping (#3852)

* stake-pool: Split up more tests to help CI, avoid warping

* Remove unnecessary updates

* Split preferred test that keeps failing

stake-pool/program/tests/deposit.rs

@@ -56,11 +56,6 @@ async fn setup(
     )
     .await;
 
-    let first_normal_slot = context.genesis_config().epoch_schedule.first_normal_slot;
-    let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
-    let mut slot = first_normal_slot;
-    context.warp_to_slot(slot).unwrap();
-
     let user = Keypair::new();
     // make stake account
     let deposit_stake = Keypair::new();
@@ -92,8 +87,8 @@ async fn setup(
     )
     .await;
 
-    slot += slots_per_epoch;
+    let first_normal_slot = context.genesis_config().epoch_schedule.first_normal_slot;
-    context.warp_to_slot(slot).unwrap();
+    context.warp_to_slot(first_normal_slot).unwrap();
     stake_pool_accounts
         .update_all(
             &mut context.banks_client,
@@ -805,226 +800,3 @@ async fn fail_with_wrong_mint_for_receiver_acc() {
         }
     }
 }
-
-#[tokio::test]
-async fn fail_with_uninitialized_validator_list() {} // TODO
-
-#[tokio::test]
-async fn fail_with_out_of_dated_pool_balances() {} // TODO
-
-#[tokio::test]
-async fn success_with_preferred_deposit() {
-    let (
-        mut context,
-        stake_pool_accounts,
-        validator_stake,
-        user,
-        deposit_stake,
-        pool_token_account,
-        _stake_lamports,
-    ) = setup(spl_token::id()).await;
-
-    stake_pool_accounts
-        .set_preferred_validator(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-            instruction::PreferredValidatorType::Deposit,
-            Some(validator_stake.vote.pubkey()),
-        )
-        .await;
-
-    let error = stake_pool_accounts
-        .deposit_stake(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-            &deposit_stake,
-            &pool_token_account,
-            &validator_stake.stake_account,
-            &user,
-        )
-        .await;
-    assert!(error.is_none());
-}
-
-#[tokio::test]
-async fn fail_with_wrong_preferred_deposit() {
-    let (
-        mut context,
-        stake_pool_accounts,
-        validator_stake,
-        user,
-        deposit_stake,
-        pool_token_account,
-        _stake_lamports,
-    ) = setup(spl_token::id()).await;
-
-    let preferred_validator = simple_add_validator_to_pool(
-        &mut context.banks_client,
-        &context.payer,
-        &context.last_blockhash,
-        &stake_pool_accounts,
-        None,
-    )
-    .await;
-
-    stake_pool_accounts
-        .set_preferred_validator(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-            instruction::PreferredValidatorType::Deposit,
-            Some(preferred_validator.vote.pubkey()),
-        )
-        .await;
-
-    let error = stake_pool_accounts
-        .deposit_stake(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-            &deposit_stake,
-            &pool_token_account,
-            &validator_stake.stake_account,
-            &user,
-        )
-        .await
-        .unwrap()
-        .unwrap();
-    match error {
-        TransactionError::InstructionError(_, InstructionError::Custom(error_index)) => {
-            assert_eq!(
-                error_index,
-                StakePoolError::IncorrectDepositVoteAddress as u32
-            );
-        }
-        _ => panic!("Wrong error occurs while try to make a deposit with wrong stake program ID"),
-    }
-}
-
-#[tokio::test]
-async fn success_with_referral_fee() {
-    let (
-        mut context,
-        stake_pool_accounts,
-        validator_stake_account,
-        user,
-        deposit_stake,
-        pool_token_account,
-        stake_lamports,
-    ) = setup(spl_token::id()).await;
-
-    let referrer = Keypair::new();
-    let referrer_token_account = Keypair::new();
-    create_token_account(
-        &mut context.banks_client,
-        &context.payer,
-        &context.last_blockhash,
-        &stake_pool_accounts.token_program_id,
-        &referrer_token_account,
-        &stake_pool_accounts.pool_mint.pubkey(),
-        &referrer,
-        &[],
-    )
-    .await
-    .unwrap();
-
-    let referrer_balance_pre =
-        get_token_balance(&mut context.banks_client, &referrer_token_account.pubkey()).await;
-
-    let mut transaction = Transaction::new_with_payer(
-        &instruction::deposit_stake(
-            &id(),
-            &stake_pool_accounts.stake_pool.pubkey(),
-            &stake_pool_accounts.validator_list.pubkey(),
-            &stake_pool_accounts.withdraw_authority,
-            &deposit_stake,
-            &user.pubkey(),
-            &validator_stake_account.stake_account,
-            &stake_pool_accounts.reserve_stake.pubkey(),
-            &pool_token_account,
-            &stake_pool_accounts.pool_fee_account.pubkey(),
-            &referrer_token_account.pubkey(),
-            &stake_pool_accounts.pool_mint.pubkey(),
-            &spl_token::id(),
-        ),
-        Some(&context.payer.pubkey()),
-    );
-    transaction.sign(&[&context.payer, &user], context.last_blockhash);
-    context
-        .banks_client
-        .process_transaction(transaction)
-        .await
-        .unwrap();
-
-    let referrer_balance_post =
-        get_token_balance(&mut context.banks_client, &referrer_token_account.pubkey()).await;
-    let stake_pool = get_account(
-        &mut context.banks_client,
-        &stake_pool_accounts.stake_pool.pubkey(),
-    )
-    .await;
-    let stake_pool =
-        try_from_slice_unchecked::<state::StakePool>(stake_pool.data.as_slice()).unwrap();
-    let rent = context.banks_client.get_rent().await.unwrap();
-    let stake_rent = rent.minimum_balance(std::mem::size_of::<stake::state::StakeState>());
-    let fee_tokens = stake_pool
-        .calc_pool_tokens_sol_deposit_fee(stake_rent)
-        .unwrap()
-        + stake_pool
-            .calc_pool_tokens_stake_deposit_fee(stake_lamports - stake_rent)
-            .unwrap();
-    let referral_fee = stake_pool_accounts.calculate_referral_fee(fee_tokens);
-    assert!(referral_fee > 0);
-    assert_eq!(referrer_balance_pre + referral_fee, referrer_balance_post);
-}
-
-#[tokio::test]
-async fn fail_with_invalid_referrer() {
-    let (
-        mut context,
-        stake_pool_accounts,
-        validator_stake_account,
-        user,
-        deposit_stake,
-        pool_token_account,
-        _stake_lamports,
-    ) = setup(spl_token::id()).await;
-
-    let invalid_token_account = Keypair::new();
-
-    let mut transaction = Transaction::new_with_payer(
-        &instruction::deposit_stake(
-            &id(),
-            &stake_pool_accounts.stake_pool.pubkey(),
-            &stake_pool_accounts.validator_list.pubkey(),
-            &stake_pool_accounts.withdraw_authority,
-            &deposit_stake,
-            &user.pubkey(),
-            &validator_stake_account.stake_account,
-            &stake_pool_accounts.reserve_stake.pubkey(),
-            &pool_token_account,
-            &stake_pool_accounts.pool_fee_account.pubkey(),
-            &invalid_token_account.pubkey(),
-            &stake_pool_accounts.pool_mint.pubkey(),
-            &spl_token::id(),
-        ),
-        Some(&context.payer.pubkey()),
-    );
-    transaction.sign(&[&context.payer, &user], context.last_blockhash);
-    let transaction_error = context
-        .banks_client
-        .process_transaction(transaction)
-        .await
-        .err()
-        .unwrap()
-        .unwrap();
-
-    match transaction_error {
-        TransactionError::InstructionError(_, InstructionError::InvalidAccountData) => (),
-        _ => panic!(
-            "Wrong error occurs while try to make a deposit with an invalid referrer account"
-        ),
-    }
-}

stake-pool/program/tests/deposit_edge_cases.rs

@@ -0,0 +1,336 @@
+#![allow(clippy::integer_arithmetic)]
+#![cfg(feature = "test-sbf")]
+
+mod helpers;
+
+use {
+    helpers::*,
+    solana_program::{
+        borsh::try_from_slice_unchecked, instruction::InstructionError, pubkey::Pubkey, stake,
+    },
+    solana_program_test::*,
+    solana_sdk::{
+        signature::{Keypair, Signer},
+        transaction::{Transaction, TransactionError},
+    },
+    spl_stake_pool::{error::StakePoolError, id, instruction, state, MINIMUM_RESERVE_LAMPORTS},
+};
+
+async fn setup(
+    token_program_id: Pubkey,
+) -> (
+    ProgramTestContext,
+    StakePoolAccounts,
+    ValidatorStakeAccount,
+    Keypair,
+    Pubkey,
+    Pubkey,
+    u64,
+) {
+    let mut context = program_test().start_with_context().await;
+
+    let stake_pool_accounts = StakePoolAccounts::new_with_token_program(token_program_id);
+    stake_pool_accounts
+        .initialize_stake_pool(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            MINIMUM_RESERVE_LAMPORTS,
+        )
+        .await
+        .unwrap();
+
+    let validator_stake_account = simple_add_validator_to_pool(
+        &mut context.banks_client,
+        &context.payer,
+        &context.last_blockhash,
+        &stake_pool_accounts,
+        None,
+    )
+    .await;
+
+    let user = Keypair::new();
+    // make stake account
+    let deposit_stake = Keypair::new();
+    let lockup = stake::state::Lockup::default();
+
+    let authorized = stake::state::Authorized {
+        staker: user.pubkey(),
+        withdrawer: user.pubkey(),
+    };
+
+    let stake_lamports = create_independent_stake_account(
+        &mut context.banks_client,
+        &context.payer,
+        &context.last_blockhash,
+        &deposit_stake,
+        &authorized,
+        &lockup,
+        TEST_STAKE_AMOUNT,
+    )
+    .await;
+
+    delegate_stake_account(
+        &mut context.banks_client,
+        &context.payer,
+        &context.last_blockhash,
+        &deposit_stake.pubkey(),
+        &user,
+        &validator_stake_account.vote.pubkey(),
+    )
+    .await;
+
+    let first_normal_slot = context.genesis_config().epoch_schedule.first_normal_slot;
+    context.warp_to_slot(first_normal_slot).unwrap();
+    stake_pool_accounts
+        .update_all(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &[validator_stake_account.vote.pubkey()],
+            false,
+        )
+        .await;
+
+    // make pool token account
+    let pool_token_account = Keypair::new();
+    create_token_account(
+        &mut context.banks_client,
+        &context.payer,
+        &context.last_blockhash,
+        &stake_pool_accounts.token_program_id,
+        &pool_token_account,
+        &stake_pool_accounts.pool_mint.pubkey(),
+        &user,
+        &[],
+    )
+    .await
+    .unwrap();
+
+    (
+        context,
+        stake_pool_accounts,
+        validator_stake_account,
+        user,
+        deposit_stake.pubkey(),
+        pool_token_account.pubkey(),
+        stake_lamports,
+    )
+}
+
+#[tokio::test]
+async fn success_with_preferred_deposit() {
+    let (
+        mut context,
+        stake_pool_accounts,
+        validator_stake,
+        user,
+        deposit_stake,
+        pool_token_account,
+        _stake_lamports,
+    ) = setup(spl_token::id()).await;
+
+    stake_pool_accounts
+        .set_preferred_validator(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            instruction::PreferredValidatorType::Deposit,
+            Some(validator_stake.vote.pubkey()),
+        )
+        .await;
+
+    let error = stake_pool_accounts
+        .deposit_stake(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &deposit_stake,
+            &pool_token_account,
+            &validator_stake.stake_account,
+            &user,
+        )
+        .await;
+    assert!(error.is_none());
+}
+
+#[tokio::test]
+async fn fail_with_wrong_preferred_deposit() {
+    let (
+        mut context,
+        stake_pool_accounts,
+        validator_stake,
+        user,
+        deposit_stake,
+        pool_token_account,
+        _stake_lamports,
+    ) = setup(spl_token::id()).await;
+
+    let preferred_validator = simple_add_validator_to_pool(
+        &mut context.banks_client,
+        &context.payer,
+        &context.last_blockhash,
+        &stake_pool_accounts,
+        None,
+    )
+    .await;
+
+    stake_pool_accounts
+        .set_preferred_validator(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            instruction::PreferredValidatorType::Deposit,
+            Some(preferred_validator.vote.pubkey()),
+        )
+        .await;
+
+    let error = stake_pool_accounts
+        .deposit_stake(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &deposit_stake,
+            &pool_token_account,
+            &validator_stake.stake_account,
+            &user,
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    match error {
+        TransactionError::InstructionError(_, InstructionError::Custom(error_index)) => {
+            assert_eq!(
+                error_index,
+                StakePoolError::IncorrectDepositVoteAddress as u32
+            );
+        }
+        _ => panic!("Wrong error occurs while try to make a deposit with wrong stake program ID"),
+    }
+}
+
+#[tokio::test]
+async fn success_with_referral_fee() {
+    let (
+        mut context,
+        stake_pool_accounts,
+        validator_stake_account,
+        user,
+        deposit_stake,
+        pool_token_account,
+        stake_lamports,
+    ) = setup(spl_token::id()).await;
+
+    let referrer = Keypair::new();
+    let referrer_token_account = Keypair::new();
+    create_token_account(
+        &mut context.banks_client,
+        &context.payer,
+        &context.last_blockhash,
+        &stake_pool_accounts.token_program_id,
+        &referrer_token_account,
+        &stake_pool_accounts.pool_mint.pubkey(),
+        &referrer,
+        &[],
+    )
+    .await
+    .unwrap();
+
+    let referrer_balance_pre =
+        get_token_balance(&mut context.banks_client, &referrer_token_account.pubkey()).await;
+
+    let mut transaction = Transaction::new_with_payer(
+        &instruction::deposit_stake(
+            &id(),
+            &stake_pool_accounts.stake_pool.pubkey(),
+            &stake_pool_accounts.validator_list.pubkey(),
+            &stake_pool_accounts.withdraw_authority,
+            &deposit_stake,
+            &user.pubkey(),
+            &validator_stake_account.stake_account,
+            &stake_pool_accounts.reserve_stake.pubkey(),
+            &pool_token_account,
+            &stake_pool_accounts.pool_fee_account.pubkey(),
+            &referrer_token_account.pubkey(),
+            &stake_pool_accounts.pool_mint.pubkey(),
+            &spl_token::id(),
+        ),
+        Some(&context.payer.pubkey()),
+    );
+    transaction.sign(&[&context.payer, &user], context.last_blockhash);
+    context
+        .banks_client
+        .process_transaction(transaction)
+        .await
+        .unwrap();
+
+    let referrer_balance_post =
+        get_token_balance(&mut context.banks_client, &referrer_token_account.pubkey()).await;
+    let stake_pool = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.stake_pool.pubkey(),
+    )
+    .await;
+    let stake_pool =
+        try_from_slice_unchecked::<state::StakePool>(stake_pool.data.as_slice()).unwrap();
+    let rent = context.banks_client.get_rent().await.unwrap();
+    let stake_rent = rent.minimum_balance(std::mem::size_of::<stake::state::StakeState>());
+    let fee_tokens = stake_pool
+        .calc_pool_tokens_sol_deposit_fee(stake_rent)
+        .unwrap()
+        + stake_pool
+            .calc_pool_tokens_stake_deposit_fee(stake_lamports - stake_rent)
+            .unwrap();
+    let referral_fee = stake_pool_accounts.calculate_referral_fee(fee_tokens);
+    assert!(referral_fee > 0);
+    assert_eq!(referrer_balance_pre + referral_fee, referrer_balance_post);
+}
+
+#[tokio::test]
+async fn fail_with_invalid_referrer() {
+    let (
+        mut context,
+        stake_pool_accounts,
+        validator_stake_account,
+        user,
+        deposit_stake,
+        pool_token_account,
+        _stake_lamports,
+    ) = setup(spl_token::id()).await;
+
+    let invalid_token_account = Keypair::new();
+
+    let mut transaction = Transaction::new_with_payer(
+        &instruction::deposit_stake(
+            &id(),
+            &stake_pool_accounts.stake_pool.pubkey(),
+            &stake_pool_accounts.validator_list.pubkey(),
+            &stake_pool_accounts.withdraw_authority,
+            &deposit_stake,
+            &user.pubkey(),
+            &validator_stake_account.stake_account,
+            &stake_pool_accounts.reserve_stake.pubkey(),
+            &pool_token_account,
+            &stake_pool_accounts.pool_fee_account.pubkey(),
+            &invalid_token_account.pubkey(),
+            &stake_pool_accounts.pool_mint.pubkey(),
+            &spl_token::id(),
+        ),
+        Some(&context.payer.pubkey()),
+    );
+    transaction.sign(&[&context.payer, &user], context.last_blockhash);
+    let transaction_error = context
+        .banks_client
+        .process_transaction(transaction)
+        .await
+        .err()
+        .unwrap()
+        .unwrap();
+
+    match transaction_error {
+        TransactionError::InstructionError(_, InstructionError::InvalidAccountData) => (),
+        _ => panic!(
+            "Wrong error occurs while try to make a deposit with an invalid referrer account"
+        ),
+    }
+}

stake-pool/program/tests/update_stake_pool_balance.rs

@@ -26,12 +26,8 @@ async fn setup(
     ProgramTestContext,
     StakePoolAccounts,
     Vec<ValidatorStakeAccount>,
-    u64,
 ) {
     let mut context = program_test().start_with_context().await;
-    let slot = context.genesis_config().epoch_schedule.first_normal_slot;
-    context.warp_to_slot(slot).unwrap();
-
     let stake_pool_accounts = StakePoolAccounts::default();
     stake_pool_accounts
         .initialize_stake_pool(
@@ -118,12 +114,12 @@ async fn setup(
         stake_accounts.push(stake_account);
     }
 
-    (context, stake_pool_accounts, stake_accounts, slot)
+    (context, stake_pool_accounts, stake_accounts)
 }
 
 #[tokio::test]
 async fn success() {
-    let (mut context, stake_pool_accounts, stake_accounts, slot) = setup(NUM_VALIDATORS).await;
+    let (mut context, stake_pool_accounts, stake_accounts) = setup(NUM_VALIDATORS).await;
 
     let pre_fee = get_token_balance(
         &mut context.banks_client,
@@ -151,15 +147,6 @@ async fn success() {
     )
     .await;
 
-    let error = stake_pool_accounts
-        .update_stake_pool_balance(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-        )
-        .await;
-    assert!(error.is_none());
-
     // Increment vote credits to earn rewards
     const VOTE_CREDITS: u64 = 1_000;
     for stake_account in &stake_accounts {
@@ -167,8 +154,8 @@ async fn success() {
     }
 
     // Update epoch
-    let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
+    let slot = context.genesis_config().epoch_schedule.first_normal_slot;
-    context.warp_to_slot(slot + slots_per_epoch).unwrap();
+    context.warp_to_slot(slot).unwrap();
 
     // Update list and pool
     let error = stake_pool_accounts
@@ -229,7 +216,7 @@ async fn success() {
 
 #[tokio::test]
 async fn success_absorbing_extra_lamports() {
-    let (mut context, stake_pool_accounts, stake_accounts, slot) = setup(NUM_VALIDATORS).await;
+    let (mut context, stake_pool_accounts, stake_accounts) = setup(NUM_VALIDATORS).await;
 
     let pre_balance = get_validator_list_sum(
         &mut context.banks_client,
@@ -251,15 +238,6 @@ async fn success_absorbing_extra_lamports() {
     )
     .await;
 
-    let error = stake_pool_accounts
-        .update_stake_pool_balance(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-        )
-        .await;
-    assert!(error.is_none());
-
     // Transfer extra funds, will be absorbed during update
     const EXTRA_STAKE_AMOUNT: u64 = 1_000_000;
     for stake_account in &stake_accounts {
@@ -277,8 +255,8 @@ async fn success_absorbing_extra_lamports() {
     let expected_fee = stake_pool.calc_epoch_fee_amount(extra_lamports).unwrap();
 
     // Update epoch
-    let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
+    let slot = context.genesis_config().epoch_schedule.first_normal_slot;
-    context.warp_to_slot(slot + slots_per_epoch).unwrap();
+    context.warp_to_slot(slot).unwrap();
 
     // Update list and pool
     let error = stake_pool_accounts

stake-pool/program/tests/update_validator_list_balance.rs

@@ -5,19 +5,10 @@ mod helpers;
 
 use {
     helpers::*,
-    solana_program::{borsh::try_from_slice_unchecked, program_pack::Pack, pubkey::Pubkey, stake},
+    solana_program::{borsh::try_from_slice_unchecked, program_pack::Pack, pubkey::Pubkey},
     solana_program_test::*,
-    solana_sdk::{
+    solana_sdk::{signature::Signer, stake::state::StakeState},
-        instruction::InstructionError,
-        signature::Signer,
-        stake::state::{Authorized, Lockup, StakeState},
-        system_instruction,
-        transaction::{Transaction, TransactionError},
-    },
     spl_stake_pool::{
-        error::StakePoolError,
-        find_stake_program_address, find_transient_stake_program_address,
-        find_withdraw_authority_program_address, id, instruction,
         state::{StakePool, StakeStatus, ValidatorList},
         MAX_VALIDATORS_TO_UPDATE, MINIMUM_RESERVE_LAMPORTS,
     },
@@ -719,371 +710,6 @@ async fn success_with_burned_tokens() {
     assert_eq!(mint.supply, stake_pool.pool_token_supply);
 }
 
-#[tokio::test]
-async fn success_ignoring_hijacked_transient_stake_with_authorized() {
-    let hijacker = Pubkey::new_unique();
-    check_ignored_hijacked_transient_stake(Some(&Authorized::auto(&hijacker)), None).await;
-}
-
-#[tokio::test]
-async fn success_ignoring_hijacked_transient_stake_with_lockup() {
-    let hijacker = Pubkey::new_unique();
-    check_ignored_hijacked_transient_stake(
-        None,
-        Some(&Lockup {
-            custodian: hijacker,
-            ..Lockup::default()
-        }),
-    )
-    .await;
-}
-
-async fn check_ignored_hijacked_transient_stake(
-    hijack_authorized: Option<&Authorized>,
-    hijack_lockup: Option<&Lockup>,
-) {
-    let num_validators = 1;
-    let (mut context, stake_pool_accounts, stake_accounts, _, lamports, _, mut slot) =
-        setup(num_validators).await;
-
-    let rent = context.banks_client.get_rent().await.unwrap();
-    let stake_rent = rent.minimum_balance(std::mem::size_of::<StakeState>());
-
-    let pre_lamports = get_validator_list_sum(
-        &mut context.banks_client,
-        &stake_pool_accounts.reserve_stake.pubkey(),
-        &stake_pool_accounts.validator_list.pubkey(),
-    )
-    .await;
-    let (withdraw_authority, _) =
-        find_withdraw_authority_program_address(&id(), &stake_pool_accounts.stake_pool.pubkey());
-
-    println!("Decrease from all validators");
-    let stake_account = &stake_accounts[0];
-    let error = stake_pool_accounts
-        .decrease_validator_stake(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-            &stake_account.stake_account,
-            &stake_account.transient_stake_account,
-            lamports,
-            stake_account.transient_stake_seed,
-        )
-        .await;
-    assert!(error.is_none());
-
-    println!("Warp one epoch so the stakes deactivate and merge");
-    let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
-    slot += slots_per_epoch;
-    context.warp_to_slot(slot).unwrap();
-
-    println!("During update, hijack the transient stake account");
-    let validator_list = stake_pool_accounts
-        .get_validator_list(&mut context.banks_client)
-        .await;
-    let transient_stake_address = find_transient_stake_program_address(
-        &id(),
-        &stake_account.vote.pubkey(),
-        &stake_pool_accounts.stake_pool.pubkey(),
-        stake_account.transient_stake_seed,
-    )
-    .0;
-    let transaction = Transaction::new_signed_with_payer(
-        &[
-            instruction::update_validator_list_balance(
-                &id(),
-                &stake_pool_accounts.stake_pool.pubkey(),
-                &stake_pool_accounts.withdraw_authority,
-                &stake_pool_accounts.validator_list.pubkey(),
-                &stake_pool_accounts.reserve_stake.pubkey(),
-                &validator_list,
-                &[stake_account.vote.pubkey()],
-                0,
-                /* no_merge = */ false,
-            ),
-            system_instruction::transfer(
-                &context.payer.pubkey(),
-                &transient_stake_address,
-                stake_rent + MINIMUM_RESERVE_LAMPORTS,
-            ),
-            stake::instruction::initialize(
-                &transient_stake_address,
-                hijack_authorized.unwrap_or(&Authorized::auto(&withdraw_authority)),
-                hijack_lockup.unwrap_or(&Lockup::default()),
-            ),
-            instruction::update_stake_pool_balance(
-                &id(),
-                &stake_pool_accounts.stake_pool.pubkey(),
-                &stake_pool_accounts.withdraw_authority,
-                &stake_pool_accounts.validator_list.pubkey(),
-                &stake_pool_accounts.reserve_stake.pubkey(),
-                &stake_pool_accounts.pool_fee_account.pubkey(),
-                &stake_pool_accounts.pool_mint.pubkey(),
-                &spl_token::id(),
-            ),
-            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());
-
-    println!("Update again normally, should be no change in the lamports");
-    stake_pool_accounts
-        .update_all(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-            stake_accounts
-                .iter()
-                .map(|v| v.vote.pubkey())
-                .collect::<Vec<Pubkey>>()
-                .as_slice(),
-            false,
-        )
-        .await;
-
-    let expected_lamports = get_validator_list_sum(
-        &mut context.banks_client,
-        &stake_pool_accounts.reserve_stake.pubkey(),
-        &stake_pool_accounts.validator_list.pubkey(),
-    )
-    .await;
-    assert_eq!(pre_lamports, expected_lamports);
-
-    let stake_pool_info = get_account(
-        &mut context.banks_client,
-        &stake_pool_accounts.stake_pool.pubkey(),
-    )
-    .await;
-    let stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool_info.data).unwrap();
-    assert_eq!(pre_lamports, stake_pool.total_lamports);
-}
-
-#[tokio::test]
-async fn success_ignoring_hijacked_validator_stake_with_authorized() {
-    let hijacker = Pubkey::new_unique();
-    check_ignored_hijacked_transient_stake(Some(&Authorized::auto(&hijacker)), None).await;
-}
-
-#[tokio::test]
-async fn success_ignoring_hijacked_validator_stake_with_lockup() {
-    let hijacker = Pubkey::new_unique();
-    check_ignored_hijacked_validator_stake(
-        None,
-        Some(&Lockup {
-            custodian: hijacker,
-            ..Lockup::default()
-        }),
-    )
-    .await;
-}
-
-async fn check_ignored_hijacked_validator_stake(
-    hijack_authorized: Option<&Authorized>,
-    hijack_lockup: Option<&Lockup>,
-) {
-    let num_validators = 1;
-    let (mut context, stake_pool_accounts, stake_accounts, _, lamports, _, mut slot) =
-        setup(num_validators).await;
-
-    let rent = context.banks_client.get_rent().await.unwrap();
-    let stake_rent = rent.minimum_balance(std::mem::size_of::<StakeState>());
-
-    let pre_lamports = get_validator_list_sum(
-        &mut context.banks_client,
-        &stake_pool_accounts.reserve_stake.pubkey(),
-        &stake_pool_accounts.validator_list.pubkey(),
-    )
-    .await;
-    let (withdraw_authority, _) =
-        find_withdraw_authority_program_address(&id(), &stake_pool_accounts.stake_pool.pubkey());
-
-    let stake_account = &stake_accounts[0];
-    let error = stake_pool_accounts
-        .decrease_validator_stake(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-            &stake_account.stake_account,
-            &stake_account.transient_stake_account,
-            lamports,
-            stake_account.transient_stake_seed,
-        )
-        .await;
-    assert!(error.is_none());
-
-    let error = stake_pool_accounts
-        .remove_validator_from_pool(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-            &stake_account.stake_account,
-            &stake_account.transient_stake_account,
-        )
-        .await;
-    assert!(error.is_none());
-
-    println!("Warp one epoch so the stakes deactivate and merge");
-    let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
-    slot += slots_per_epoch;
-    context.warp_to_slot(slot).unwrap();
-
-    println!("During update, hijack the validator stake account");
-    let validator_list = stake_pool_accounts
-        .get_validator_list(&mut context.banks_client)
-        .await;
-    let transaction = Transaction::new_signed_with_payer(
-        &[
-            instruction::update_validator_list_balance(
-                &id(),
-                &stake_pool_accounts.stake_pool.pubkey(),
-                &stake_pool_accounts.withdraw_authority,
-                &stake_pool_accounts.validator_list.pubkey(),
-                &stake_pool_accounts.reserve_stake.pubkey(),
-                &validator_list,
-                &[stake_account.vote.pubkey()],
-                0,
-                /* no_merge = */ false,
-            ),
-            system_instruction::transfer(
-                &context.payer.pubkey(),
-                &stake_account.stake_account,
-                stake_rent + MINIMUM_RESERVE_LAMPORTS,
-            ),
-            stake::instruction::initialize(
-                &stake_account.stake_account,
-                hijack_authorized.unwrap_or(&Authorized::auto(&withdraw_authority)),
-                hijack_lockup.unwrap_or(&Lockup::default()),
-            ),
-            instruction::update_stake_pool_balance(
-                &id(),
-                &stake_pool_accounts.stake_pool.pubkey(),
-                &stake_pool_accounts.withdraw_authority,
-                &stake_pool_accounts.validator_list.pubkey(),
-                &stake_pool_accounts.reserve_stake.pubkey(),
-                &stake_pool_accounts.pool_fee_account.pubkey(),
-                &stake_pool_accounts.pool_mint.pubkey(),
-                &spl_token::id(),
-            ),
-            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());
-
-    println!("Update again normally, should be no change in the lamports");
-    stake_pool_accounts
-        .update_all(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-            stake_accounts
-                .iter()
-                .map(|v| v.vote.pubkey())
-                .collect::<Vec<Pubkey>>()
-                .as_slice(),
-            false,
-        )
-        .await;
-
-    let expected_lamports = get_validator_list_sum(
-        &mut context.banks_client,
-        &stake_pool_accounts.reserve_stake.pubkey(),
-        &stake_pool_accounts.validator_list.pubkey(),
-    )
-    .await;
-    assert_eq!(pre_lamports, expected_lamports);
-
-    let stake_pool_info = get_account(
-        &mut context.banks_client,
-        &stake_pool_accounts.stake_pool.pubkey(),
-    )
-    .await;
-    let stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool_info.data).unwrap();
-    assert_eq!(pre_lamports, stake_pool.total_lamports);
-
-    println!("Fail adding validator back in with first seed");
-    let error = stake_pool_accounts
-        .add_validator_to_pool(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-            &stake_account.stake_account,
-            &stake_account.vote.pubkey(),
-            stake_account.validator_stake_seed,
-        )
-        .await
-        .unwrap()
-        .unwrap();
-    assert_eq!(
-        error,
-        TransactionError::InstructionError(
-            0,
-            InstructionError::Custom(StakePoolError::AlreadyInUse as u32),
-        )
-    );
-
-    println!("Succeed adding validator back in with new seed");
-    let seed = NonZeroU32::new(1);
-    let validator = stake_account.vote.pubkey();
-    let (stake_account, _) = find_stake_program_address(
-        &id(),
-        &validator,
-        &stake_pool_accounts.stake_pool.pubkey(),
-        seed,
-    );
-    let error = stake_pool_accounts
-        .add_validator_to_pool(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-            &stake_account,
-            &validator,
-            seed,
-        )
-        .await;
-    assert!(error.is_none());
-
-    let stake_pool_info = get_account(
-        &mut context.banks_client,
-        &stake_pool_accounts.stake_pool.pubkey(),
-    )
-    .await;
-    let stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool_info.data).unwrap();
-    assert_eq!(pre_lamports, stake_pool.total_lamports);
-
-    let expected_lamports = get_validator_list_sum(
-        &mut context.banks_client,
-        &stake_pool_accounts.reserve_stake.pubkey(),
-        &stake_pool_accounts.validator_list.pubkey(),
-    )
-    .await;
-    assert_eq!(pre_lamports, expected_lamports);
-}
-
 #[tokio::test]
 async fn fail_with_uninitialized_validator_list() {} // TODO
 

stake-pool/program/tests/update_validator_list_balance_hijack.rs

@@ -0,0 +1,521 @@
+#![allow(clippy::integer_arithmetic)]
+#![cfg(feature = "test-sbf")]
+
+mod helpers;
+
+use {
+    helpers::*,
+    solana_program::{borsh::try_from_slice_unchecked, pubkey::Pubkey, stake},
+    solana_program_test::*,
+    solana_sdk::{
+        instruction::InstructionError,
+        signature::Signer,
+        stake::state::{Authorized, Lockup, StakeState},
+        system_instruction,
+        transaction::{Transaction, TransactionError},
+    },
+    spl_stake_pool::{
+        error::StakePoolError, find_stake_program_address, find_transient_stake_program_address,
+        find_withdraw_authority_program_address, id, instruction, state::StakePool,
+        MINIMUM_RESERVE_LAMPORTS,
+    },
+    std::num::NonZeroU32,
+};
+
+async fn setup(
+    num_validators: usize,
+) -> (
+    ProgramTestContext,
+    StakePoolAccounts,
+    Vec<ValidatorStakeAccount>,
+    Vec<DepositStakeAccount>,
+    u64,
+    u64,
+    u64,
+) {
+    let mut context = program_test().start_with_context().await;
+    let first_normal_slot = context.genesis_config().epoch_schedule.first_normal_slot;
+    let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
+    let mut slot = first_normal_slot;
+    context.warp_to_slot(slot).unwrap();
+
+    let reserve_stake_amount = TEST_STAKE_AMOUNT * 2 * num_validators as u64;
+    let stake_pool_accounts = StakePoolAccounts::default();
+    stake_pool_accounts
+        .initialize_stake_pool(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            reserve_stake_amount + MINIMUM_RESERVE_LAMPORTS,
+        )
+        .await
+        .unwrap();
+
+    // Add several accounts with some stake
+    let mut stake_accounts: Vec<ValidatorStakeAccount> = vec![];
+    let mut deposit_accounts: Vec<DepositStakeAccount> = vec![];
+    for i in 0..num_validators {
+        let stake_account = ValidatorStakeAccount::new(
+            &stake_pool_accounts.stake_pool.pubkey(),
+            NonZeroU32::new(i as u32),
+            u64::MAX,
+        );
+        create_vote(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &stake_account.validator,
+            &stake_account.vote,
+        )
+        .await;
+
+        let error = stake_pool_accounts
+            .add_validator_to_pool(
+                &mut context.banks_client,
+                &context.payer,
+                &context.last_blockhash,
+                &stake_account.stake_account,
+                &stake_account.vote.pubkey(),
+                stake_account.validator_stake_seed,
+            )
+            .await;
+        assert!(error.is_none());
+
+        let deposit_account = DepositStakeAccount::new_with_vote(
+            stake_account.vote.pubkey(),
+            stake_account.stake_account,
+            TEST_STAKE_AMOUNT,
+        );
+        deposit_account
+            .create_and_delegate(
+                &mut context.banks_client,
+                &context.payer,
+                &context.last_blockhash,
+            )
+            .await;
+
+        stake_accounts.push(stake_account);
+        deposit_accounts.push(deposit_account);
+    }
+
+    // Warp forward so the stakes properly activate, and deposit
+    slot += slots_per_epoch;
+    context.warp_to_slot(slot).unwrap();
+
+    stake_pool_accounts
+        .update_all(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            stake_accounts
+                .iter()
+                .map(|v| v.vote.pubkey())
+                .collect::<Vec<Pubkey>>()
+                .as_slice(),
+            false,
+        )
+        .await;
+
+    for deposit_account in &mut deposit_accounts {
+        deposit_account
+            .deposit_stake(
+                &mut context.banks_client,
+                &context.payer,
+                &context.last_blockhash,
+                &stake_pool_accounts,
+            )
+            .await;
+    }
+
+    slot += slots_per_epoch;
+    context.warp_to_slot(slot).unwrap();
+
+    stake_pool_accounts
+        .update_all(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            stake_accounts
+                .iter()
+                .map(|v| v.vote.pubkey())
+                .collect::<Vec<Pubkey>>()
+                .as_slice(),
+            false,
+        )
+        .await;
+
+    (
+        context,
+        stake_pool_accounts,
+        stake_accounts,
+        deposit_accounts,
+        TEST_STAKE_AMOUNT,
+        reserve_stake_amount,
+        slot,
+    )
+}
+
+#[tokio::test]
+async fn success_ignoring_hijacked_transient_stake_with_authorized() {
+    let hijacker = Pubkey::new_unique();
+    check_ignored_hijacked_transient_stake(Some(&Authorized::auto(&hijacker)), None).await;
+}
+
+#[tokio::test]
+async fn success_ignoring_hijacked_transient_stake_with_lockup() {
+    let hijacker = Pubkey::new_unique();
+    check_ignored_hijacked_transient_stake(
+        None,
+        Some(&Lockup {
+            custodian: hijacker,
+            ..Lockup::default()
+        }),
+    )
+    .await;
+}
+
+async fn check_ignored_hijacked_transient_stake(
+    hijack_authorized: Option<&Authorized>,
+    hijack_lockup: Option<&Lockup>,
+) {
+    let num_validators = 1;
+    let (mut context, stake_pool_accounts, stake_accounts, _, lamports, _, mut slot) =
+        setup(num_validators).await;
+
+    let rent = context.banks_client.get_rent().await.unwrap();
+    let stake_rent = rent.minimum_balance(std::mem::size_of::<StakeState>());
+
+    let pre_lamports = get_validator_list_sum(
+        &mut context.banks_client,
+        &stake_pool_accounts.reserve_stake.pubkey(),
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let (withdraw_authority, _) =
+        find_withdraw_authority_program_address(&id(), &stake_pool_accounts.stake_pool.pubkey());
+
+    println!("Decrease from all validators");
+    let stake_account = &stake_accounts[0];
+    let error = stake_pool_accounts
+        .decrease_validator_stake(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &stake_account.stake_account,
+            &stake_account.transient_stake_account,
+            lamports,
+            stake_account.transient_stake_seed,
+        )
+        .await;
+    assert!(error.is_none());
+
+    println!("Warp one epoch so the stakes deactivate and merge");
+    let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
+    slot += slots_per_epoch;
+    context.warp_to_slot(slot).unwrap();
+
+    println!("During update, hijack the transient stake account");
+    let validator_list = stake_pool_accounts
+        .get_validator_list(&mut context.banks_client)
+        .await;
+    let transient_stake_address = find_transient_stake_program_address(
+        &id(),
+        &stake_account.vote.pubkey(),
+        &stake_pool_accounts.stake_pool.pubkey(),
+        stake_account.transient_stake_seed,
+    )
+    .0;
+    let transaction = Transaction::new_signed_with_payer(
+        &[
+            instruction::update_validator_list_balance(
+                &id(),
+                &stake_pool_accounts.stake_pool.pubkey(),
+                &stake_pool_accounts.withdraw_authority,
+                &stake_pool_accounts.validator_list.pubkey(),
+                &stake_pool_accounts.reserve_stake.pubkey(),
+                &validator_list,
+                &[stake_account.vote.pubkey()],
+                0,
+                /* no_merge = */ false,
+            ),
+            system_instruction::transfer(
+                &context.payer.pubkey(),
+                &transient_stake_address,
+                stake_rent + MINIMUM_RESERVE_LAMPORTS,
+            ),
+            stake::instruction::initialize(
+                &transient_stake_address,
+                hijack_authorized.unwrap_or(&Authorized::auto(&withdraw_authority)),
+                hijack_lockup.unwrap_or(&Lockup::default()),
+            ),
+            instruction::update_stake_pool_balance(
+                &id(),
+                &stake_pool_accounts.stake_pool.pubkey(),
+                &stake_pool_accounts.withdraw_authority,
+                &stake_pool_accounts.validator_list.pubkey(),
+                &stake_pool_accounts.reserve_stake.pubkey(),
+                &stake_pool_accounts.pool_fee_account.pubkey(),
+                &stake_pool_accounts.pool_mint.pubkey(),
+                &spl_token::id(),
+            ),
+            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());
+
+    println!("Update again normally, should be no change in the lamports");
+    stake_pool_accounts
+        .update_all(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            stake_accounts
+                .iter()
+                .map(|v| v.vote.pubkey())
+                .collect::<Vec<Pubkey>>()
+                .as_slice(),
+            false,
+        )
+        .await;
+
+    let expected_lamports = get_validator_list_sum(
+        &mut context.banks_client,
+        &stake_pool_accounts.reserve_stake.pubkey(),
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    assert_eq!(pre_lamports, expected_lamports);
+
+    let stake_pool_info = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.stake_pool.pubkey(),
+    )
+    .await;
+    let stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool_info.data).unwrap();
+    assert_eq!(pre_lamports, stake_pool.total_lamports);
+}
+
+#[tokio::test]
+async fn success_ignoring_hijacked_validator_stake_with_authorized() {
+    let hijacker = Pubkey::new_unique();
+    check_ignored_hijacked_transient_stake(Some(&Authorized::auto(&hijacker)), None).await;
+}
+
+#[tokio::test]
+async fn success_ignoring_hijacked_validator_stake_with_lockup() {
+    let hijacker = Pubkey::new_unique();
+    check_ignored_hijacked_validator_stake(
+        None,
+        Some(&Lockup {
+            custodian: hijacker,
+            ..Lockup::default()
+        }),
+    )
+    .await;
+}
+
+async fn check_ignored_hijacked_validator_stake(
+    hijack_authorized: Option<&Authorized>,
+    hijack_lockup: Option<&Lockup>,
+) {
+    let num_validators = 1;
+    let (mut context, stake_pool_accounts, stake_accounts, _, lamports, _, mut slot) =
+        setup(num_validators).await;
+
+    let rent = context.banks_client.get_rent().await.unwrap();
+    let stake_rent = rent.minimum_balance(std::mem::size_of::<StakeState>());
+
+    let pre_lamports = get_validator_list_sum(
+        &mut context.banks_client,
+        &stake_pool_accounts.reserve_stake.pubkey(),
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let (withdraw_authority, _) =
+        find_withdraw_authority_program_address(&id(), &stake_pool_accounts.stake_pool.pubkey());
+
+    let stake_account = &stake_accounts[0];
+    let error = stake_pool_accounts
+        .decrease_validator_stake(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &stake_account.stake_account,
+            &stake_account.transient_stake_account,
+            lamports,
+            stake_account.transient_stake_seed,
+        )
+        .await;
+    assert!(error.is_none());
+
+    let error = stake_pool_accounts
+        .remove_validator_from_pool(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &stake_account.stake_account,
+            &stake_account.transient_stake_account,
+        )
+        .await;
+    assert!(error.is_none());
+
+    println!("Warp one epoch so the stakes deactivate and merge");
+    let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
+    slot += slots_per_epoch;
+    context.warp_to_slot(slot).unwrap();
+
+    println!("During update, hijack the validator stake account");
+    let validator_list = stake_pool_accounts
+        .get_validator_list(&mut context.banks_client)
+        .await;
+    let transaction = Transaction::new_signed_with_payer(
+        &[
+            instruction::update_validator_list_balance(
+                &id(),
+                &stake_pool_accounts.stake_pool.pubkey(),
+                &stake_pool_accounts.withdraw_authority,
+                &stake_pool_accounts.validator_list.pubkey(),
+                &stake_pool_accounts.reserve_stake.pubkey(),
+                &validator_list,
+                &[stake_account.vote.pubkey()],
+                0,
+                /* no_merge = */ false,
+            ),
+            system_instruction::transfer(
+                &context.payer.pubkey(),
+                &stake_account.stake_account,
+                stake_rent + MINIMUM_RESERVE_LAMPORTS,
+            ),
+            stake::instruction::initialize(
+                &stake_account.stake_account,
+                hijack_authorized.unwrap_or(&Authorized::auto(&withdraw_authority)),
+                hijack_lockup.unwrap_or(&Lockup::default()),
+            ),
+            instruction::update_stake_pool_balance(
+                &id(),
+                &stake_pool_accounts.stake_pool.pubkey(),
+                &stake_pool_accounts.withdraw_authority,
+                &stake_pool_accounts.validator_list.pubkey(),
+                &stake_pool_accounts.reserve_stake.pubkey(),
+                &stake_pool_accounts.pool_fee_account.pubkey(),
+                &stake_pool_accounts.pool_mint.pubkey(),
+                &spl_token::id(),
+            ),
+            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());
+
+    println!("Update again normally, should be no change in the lamports");
+    stake_pool_accounts
+        .update_all(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            stake_accounts
+                .iter()
+                .map(|v| v.vote.pubkey())
+                .collect::<Vec<Pubkey>>()
+                .as_slice(),
+            false,
+        )
+        .await;
+
+    let expected_lamports = get_validator_list_sum(
+        &mut context.banks_client,
+        &stake_pool_accounts.reserve_stake.pubkey(),
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    assert_eq!(pre_lamports, expected_lamports);
+
+    let stake_pool_info = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.stake_pool.pubkey(),
+    )
+    .await;
+    let stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool_info.data).unwrap();
+    assert_eq!(pre_lamports, stake_pool.total_lamports);
+
+    println!("Fail adding validator back in with first seed");
+    let error = stake_pool_accounts
+        .add_validator_to_pool(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &stake_account.stake_account,
+            &stake_account.vote.pubkey(),
+            stake_account.validator_stake_seed,
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(
+        error,
+        TransactionError::InstructionError(
+            0,
+            InstructionError::Custom(StakePoolError::AlreadyInUse as u32),
+        )
+    );
+
+    println!("Succeed adding validator back in with new seed");
+    let seed = NonZeroU32::new(1);
+    let validator = stake_account.vote.pubkey();
+    let (stake_account, _) = find_stake_program_address(
+        &id(),
+        &validator,
+        &stake_pool_accounts.stake_pool.pubkey(),
+        seed,
+    );
+    let error = stake_pool_accounts
+        .add_validator_to_pool(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &stake_account,
+            &validator,
+            seed,
+        )
+        .await;
+    assert!(error.is_none());
+
+    let stake_pool_info = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.stake_pool.pubkey(),
+    )
+    .await;
+    let stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool_info.data).unwrap();
+    assert_eq!(pre_lamports, stake_pool.total_lamports);
+
+    let expected_lamports = get_validator_list_sum(
+        &mut context.banks_client,
+        &stake_pool_accounts.reserve_stake.pubkey(),
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    assert_eq!(pre_lamports, expected_lamports);
+}

stake-pool/program/tests/withdraw_edge_cases.rs

@@ -474,7 +474,7 @@ async fn success_with_reserve() {
 }
 
 #[tokio::test]
-async fn success_and_fail_with_preferred_withdraw() {
+async fn success_with_empty_preferred_withdraw() {
     let (
         mut context,
         stake_pool_accounts,
@@ -520,17 +520,39 @@ async fn success_and_fail_with_preferred_withdraw() {
         )
         .await;
     assert!(error.is_none());
+}
 
-    // Deposit into preferred, then fail
+#[tokio::test]
-    let user_stake_recipient = Keypair::new();
+async fn success_and_fail_with_preferred_withdraw() {
-    create_blank_stake_account(
+    let (
+        mut context,
+        stake_pool_accounts,
+        validator_stake,
+        deposit_info,
+        user_transfer_authority,
+        user_stake_recipient,
+        tokens_to_burn,
+    ) = setup(spl_token::id()).await;
+
+    let preferred_validator = simple_add_validator_to_pool(
         &mut context.banks_client,
         &context.payer,
         &context.last_blockhash,
-        &user_stake_recipient,
+        &stake_pool_accounts,
+        None,
     )
     .await;
 
+    stake_pool_accounts
+        .set_preferred_validator(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            instruction::PreferredValidatorType::Withdraw,
+            Some(preferred_validator.vote.pubkey()),
+        )
+        .await;
+
     let _preferred_deposit = simple_deposit_stake(
         &mut context.banks_client,
         &context.payer,
@@ -542,6 +564,7 @@ async fn success_and_fail_with_preferred_withdraw() {
     .await
     .unwrap();
 
+    let new_authority = Pubkey::new_unique();
     let error = stake_pool_accounts
         .withdraw_stake(
             &mut context.banks_client,