stake-pool: Add ability to remove a validator that has deactivating transient stake (#1624)

* Add status enum

* Add ability to remove validator with transient stake

* Only account validator stake if active

* Fix merge conflicts

stake-pool/program/src/instruction.rs

@@ -519,7 +519,7 @@ pub fn update_stake_pool_balance(
     let accounts = vec![
         AccountMeta::new(*stake_pool, false),
         AccountMeta::new_readonly(*withdraw_authority, false),
-        AccountMeta::new_readonly(*validator_list_storage, false),
+        AccountMeta::new(*validator_list_storage, false),
         AccountMeta::new_readonly(*reserve_stake, false),
         AccountMeta::new(*manager_fee_account, false),
         AccountMeta::new(*stake_pool_mint, false),

stake-pool/program/src/processor.rs

@@ -7,7 +7,7 @@ use {
         find_deposit_authority_program_address,
         instruction::StakePoolInstruction,
         minimum_reserve_lamports, minimum_stake_lamports, stake_program,
-        state::{AccountType, Fee, StakePool, ValidatorList, ValidatorStakeInfo},
+        state::{AccountType, Fee, StakePool, StakeStatus, ValidatorList, ValidatorStakeInfo},
         AUTHORITY_DEPOSIT, AUTHORITY_WITHDRAW, MINIMUM_ACTIVE_STAKE, TRANSIENT_STAKE_SEED,
     },
     borsh::{BorshDeserialize, BorshSerialize},
@@ -15,7 +15,7 @@ use {
     solana_program::{
         account_info::next_account_info,
         account_info::AccountInfo,
-        clock::Clock,
+        clock::{Clock, Epoch},
         decode_error::DecodeError,
         entrypoint::ProgramResult,
         msg,
@@ -747,6 +747,7 @@ impl Processor {
         )?;
 
         validator_list.validators.push(ValidatorStakeInfo {
+            status: StakeStatus::Active,
             vote_account_address,
             stake_lamports: stake_lamports.saturating_sub(minimum_lamport_amount),
             last_update_epoch: clock.epoch,
@@ -814,20 +815,16 @@ impl Processor {
             transient_stake_account_info.key,
             &vote_account_address,
         )?;
-        // check that the transient stake account doesn't exist
+
-        if get_stake_state(transient_stake_account_info).is_ok() {
+        let maybe_validator_list_entry = validator_list.find_mut(&vote_account_address);
+        if maybe_validator_list_entry.is_none() {
             msg!(
-                "Transient stake {} exists, can't remove stake {} on validator {}",
+                "Vote account {} not found in stake pool",
-                transient_stake_account_info.key,
-                stake_account_info.key,
                 vote_account_address
             );
-            return Err(StakePoolError::WrongStakeState.into());
-        }
-
-        if !validator_list.contains(&vote_account_address) {
             return Err(StakePoolError::ValidatorNotFound.into());
         }
+        let mut validator_list_entry = maybe_validator_list_entry.unwrap();
 
         let stake_lamports = **stake_account_info.lamports.borrow();
         let required_lamports = minimum_stake_lamports(&meta);
@@ -840,6 +837,24 @@ impl Processor {
             return Err(StakePoolError::StakeLamportsNotEqualToMinimum.into());
         }
 
+        // check that the transient stake account doesn't exist
+        let new_status = if let Ok((_meta, stake)) = get_stake_state(transient_stake_account_info) {
+            if stake.delegation.deactivation_epoch == Epoch::MAX {
+                msg!(
+                    "Transient stake {} activating, can't remove stake {} on validator {}",
+                    transient_stake_account_info.key,
+                    stake_account_info.key,
+                    vote_account_address
+                );
+                return Err(StakePoolError::WrongStakeState.into());
+            } else {
+                // stake is deactivating, mark the entry as such
+                StakeStatus::DeactivatingTransient
+            }
+        } else {
+            StakeStatus::ReadyForRemoval
+        };
+
         Self::stake_authorize_signed(
             stake_pool_info.key,
             stake_account_info.clone(),
@@ -851,9 +866,13 @@ impl Processor {
             stake_program_info.clone(),
         )?;
 
-        validator_list
+        match new_status {
-            .validators
+            StakeStatus::DeactivatingTransient => validator_list_entry.status = new_status,
-            .retain(|item| item.vote_account_address != vote_account_address);
+            StakeStatus::ReadyForRemoval => validator_list
+                .validators
+                .retain(|item| item.vote_account_address != vote_account_address),
+            _ => unreachable!(),
+        }
         validator_list.serialize(&mut *validator_list_info.data.borrow_mut())?;
 
         Ok(())
@@ -1239,6 +1258,11 @@ impl Processor {
                             stake_history_info.clone(),
                             stake_program_info.clone(),
                         )?;
+                        if validator_stake_record.status == StakeStatus::DeactivatingTransient {
+                            // the validator stake was previously removed, and
+                            // now this entry can be removed totally
+                            validator_stake_record.status = StakeStatus::ReadyForRemoval;
+                        }
                     }
                 }
                 Some(stake_program::StakeState::Stake(_, stake)) => {
@@ -1257,6 +1281,11 @@ impl Processor {
                             stake_history_info.clone(),
                             stake_program_info.clone(),
                         )?;
+                        if validator_stake_record.status == StakeStatus::DeactivatingTransient {
+                            // the validator stake was previously removed, and
+                            // now this entry can be removed totally
+                            validator_stake_record.status = StakeStatus::ReadyForRemoval;
+                        }
                     } else if stake.delegation.activation_epoch < clock.epoch {
                         if let Some(stake_program::StakeState::Stake(_, validator_stake)) =
                             validator_stake_state
@@ -1298,15 +1327,19 @@ impl Processor {
             //  * any other state / not a stake -> error state, but account for transient stake
             match validator_stake_state {
                 Some(stake_program::StakeState::Stake(meta, _)) => {
-                    stake_lamports += validator_stake_info
+                    if validator_stake_record.status == StakeStatus::Active {
-                        .lamports()
+                        stake_lamports += validator_stake_info
-                        .saturating_sub(minimum_stake_lamports(&meta));
+                            .lamports()
+                            .saturating_sub(minimum_stake_lamports(&meta));
+                    } else {
+                        msg!("Validator stake account no longer part of the pool, ignoring");
+                    }
                 }
                 Some(stake_program::StakeState::Initialized(_))
                 | Some(stake_program::StakeState::Uninitialized)
                 | Some(stake_program::StakeState::RewardsPool)
                 | None => {
-                    msg!("Validator stake account no longer part of the pool, not considering");
+                    msg!("Validator stake account no longer part of the pool, ignoring");
                 }
             }
 
@@ -1355,7 +1388,7 @@ impl Processor {
             return Err(ProgramError::IncorrectProgramId);
         }
 
-        let validator_list =
+        let mut validator_list =
             try_from_slice_unchecked::<ValidatorList>(&validator_list_info.data.borrow())?;
         if !validator_list.is_valid() {
             return Err(StakePoolError::InvalidState.into());
@@ -1375,7 +1408,7 @@ impl Processor {
                 msg!("Reserve stake account in unknown state, aborting");
                 return Err(StakePoolError::WrongStakeState.into());
             };
-        for validator_stake_record in validator_list.validators {
+        for validator_stake_record in &validator_list.validators {
             if validator_stake_record.last_update_epoch < clock.epoch {
                 return Err(StakePoolError::StakeListOutOfDate.into());
             }
@@ -1406,6 +1439,10 @@ impl Processor {
                 .checked_add(fee)
                 .ok_or(StakePoolError::CalculationFailure)?;
         }
+        validator_list
+            .validators
+            .retain(|item| item.status != StakeStatus::ReadyForRemoval);
+        validator_list.serialize(&mut *validator_list_info.data.borrow_mut())?;
         stake_pool.total_stake_lamports = total_stake_lamports;
         stake_pool.last_update_epoch = clock.epoch;
         stake_pool.serialize(&mut *stake_pool_info.data.borrow_mut())?;
@@ -1507,6 +1544,11 @@ impl Processor {
             .find_mut(&vote_account_address)
             .ok_or(StakePoolError::ValidatorNotFound)?;
 
+        if validator_list_item.status != StakeStatus::Active {
+            msg!("Validator is marked for removal and no longer accepting deposits");
+            return Err(StakePoolError::ValidatorNotFound.into());
+        }
+
         let stake_lamports = **stake_info.lamports.borrow();
         let new_pool_tokens = stake_pool
             .calc_pool_tokens_for_deposit(stake_lamports)
@@ -1684,6 +1726,11 @@ impl Processor {
                 .find_mut(&vote_account_address)
                 .ok_or(StakePoolError::ValidatorNotFound)?;
 
+            if validator_list_item.status != StakeStatus::Active {
+                msg!("Validator is marked for removal and no longer allowing withdrawals");
+                return Err(StakePoolError::ValidatorNotFound.into());
+            }
+
             let required_lamports = minimum_stake_lamports(&meta);
             let current_lamports = stake_split_from.lamports();
             let remaining_lamports = current_lamports.saturating_sub(withdraw_lamports);

stake-pool/program/src/state.rs

@@ -285,10 +285,32 @@ pub struct ValidatorList {
     pub validators: Vec<ValidatorStakeInfo>,
 }
 
+/// Status of the stake account in the validator list, for accounting
+#[derive(Copy, Clone, Debug, PartialEq, BorshDeserialize, BorshSerialize, BorshSchema)]
+pub enum StakeStatus {
+    /// Stake account is active, there may be a transient stake as well
+    Active,
+    /// Only transient stake account exists, when a transient stake is
+    /// deactivating during validator removal
+    DeactivatingTransient,
+    /// No more validator stake accounts exist, entry ready for removal during
+    /// `UpdateStakePoolBalance`
+    ReadyForRemoval,
+}
+
+impl Default for StakeStatus {
+    fn default() -> Self {
+        Self::Active
+    }
+}
+
 /// Information about the singe validator stake account
 #[repr(C)]
 #[derive(Clone, Copy, Debug, Default, PartialEq, BorshDeserialize, BorshSerialize, BorshSchema)]
 pub struct ValidatorStakeInfo {
+    /// Status of the validator stake account
+    pub status: StakeStatus,
+
     /// Validator vote account address
     pub vote_account_address: Pubkey,
 
@@ -314,7 +336,7 @@ impl ValidatorList {
     /// Calculate the number of validator entries that fit in the provided length
     pub fn calculate_max_validators(buffer_length: usize) -> usize {
         let header_size = 1 + 4 + 4;
-        buffer_length.saturating_sub(header_size) / 48
+        buffer_length.saturating_sub(header_size) / 49
     }
 
     /// Check if contains validator with particular pubkey
@@ -402,16 +424,19 @@ mod test {
             max_validators,
             validators: vec![
                 ValidatorStakeInfo {
+                    status: StakeStatus::Active,
                     vote_account_address: Pubkey::new_from_array([1; 32]),
                     stake_lamports: 123456789,
                     last_update_epoch: 987654321,
                 },
                 ValidatorStakeInfo {
+                    status: StakeStatus::DeactivatingTransient,
                     vote_account_address: Pubkey::new_from_array([2; 32]),
                     stake_lamports: 998877665544,
                     last_update_epoch: 11223445566,
                 },
                 ValidatorStakeInfo {
+                    status: StakeStatus::ReadyForRemoval,
                     vote_account_address: Pubkey::new_from_array([3; 32]),
                     stake_lamports: 0,
                     last_update_epoch: 999999999999999,

stake-pool/program/tests/decrease.rs

@@ -49,7 +49,8 @@ async fn setup() -> (
         &validator_stake_account,
         100_000_000,
     )
-    .await;
+    .await
+    .unwrap();
 
     let lamports = deposit_info.stake_lamports / 2;
 

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

@@ -1010,7 +1010,7 @@ pub async fn simple_deposit(
     stake_pool_accounts: &StakePoolAccounts,
     validator_stake_account: &ValidatorStakeAccount,
     stake_lamports: u64,
-) -> DepositStakeAccount {
+) -> Option<DepositStakeAccount> {
     let authority = Keypair::new();
     // make stake account
     let stake = Keypair::new();
@@ -1064,11 +1064,11 @@ pub async fn simple_deposit(
             &authority,
         )
         .await
-        .unwrap();
+        .ok()?;
 
     let pool_tokens = get_token_balance(banks_client, &pool_account.pubkey()).await;
 
-    DepositStakeAccount {
+    Some(DepositStakeAccount {
         authority,
         stake,
         pool_account,
@@ -1076,7 +1076,7 @@ pub async fn simple_deposit(
         pool_tokens,
         vote_account,
         validator_stake_account,
-    }
+    })
 }
 
 pub async fn get_validator_list_sum(

stake-pool/program/tests/increase.rs

@@ -54,7 +54,8 @@ async fn setup() -> (
         &validator_stake_account,
         5_000_000,
     )
-    .await;
+    .await
+    .unwrap();
 
     (
         banks_client,

stake-pool/program/tests/update_stake_pool_balance.rs

@@ -51,7 +51,8 @@ async fn setup() -> (
             &validator_stake_account,
             TEST_STAKE_AMOUNT,
         )
-        .await;
+        .await
+        .unwrap();
 
         stake_accounts.push(validator_stake_account);
     }

stake-pool/program/tests/update_validator_list_balance.rs

@@ -9,7 +9,10 @@ use {
     solana_program_test::*,
     solana_sdk::signature::{Keypair, Signer},
     spl_stake_pool::{
-        stake_program, state::StakePool, MAX_VALIDATORS_TO_UPDATE, MINIMUM_ACTIVE_STAKE,
+        borsh::try_from_slice_unchecked,
+        stake_program,
+        state::{StakePool, StakeStatus, ValidatorList},
+        MAX_VALIDATORS_TO_UPDATE, MINIMUM_ACTIVE_STAKE,
     },
 };
 
@@ -474,7 +477,139 @@ async fn merge_into_validator_stake() {
 }
 
 #[tokio::test]
-async fn max_validators() {}
+#[ignore]
+async fn merge_transient_stake_after_remove() {
+    let (mut context, stake_pool_accounts, stake_accounts, lamports, reserve_lamports, mut slot) =
+        setup(1).await;
+
+    let rent = context.banks_client.get_rent().await.unwrap();
+    let stake_rent = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
+    let deactivated_lamports = lamports + stake_rent;
+    let new_authority = Pubkey::new_unique();
+    // Decrease and remove all validators
+    for stake_account in &stake_accounts {
+        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,
+                deactivated_lamports,
+            )
+            .await;
+        assert!(error.is_none());
+        let error = stake_pool_accounts
+            .remove_validator_from_pool(
+                &mut context.banks_client,
+                &context.payer,
+                &context.last_blockhash,
+                &new_authority,
+                &stake_account.stake_account,
+                &stake_account.transient_stake_account,
+            )
+            .await;
+        assert!(error.is_none());
+    }
+
+    // Warp forward to merge time
+    let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
+    slot += slots_per_epoch;
+    context.warp_to_slot(slot).unwrap();
+
+    // Update without merge, status should be DeactivatingTransient
+    let error = 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(),
+            true,
+        )
+        .await;
+    assert!(error.is_none());
+
+    let validator_list = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let validator_list =
+        try_from_slice_unchecked::<ValidatorList>(validator_list.data.as_slice()).unwrap();
+    assert_eq!(validator_list.validators.len(), 1);
+    assert_eq!(
+        validator_list.validators[0].status,
+        StakeStatus::DeactivatingTransient
+    );
+    assert_eq!(
+        validator_list.validators[0].stake_lamports,
+        deactivated_lamports
+    );
+
+    // Update with merge, status should be ReadyForRemoval and no lamports
+    let error = stake_pool_accounts
+        .update_validator_list_balance(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            stake_accounts
+                .iter()
+                .map(|v| v.vote.pubkey())
+                .collect::<Vec<Pubkey>>()
+                .as_slice(),
+            false,
+        )
+        .await;
+    assert!(error.is_none());
+
+    let validator_list = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let validator_list =
+        try_from_slice_unchecked::<ValidatorList>(validator_list.data.as_slice()).unwrap();
+    assert_eq!(validator_list.validators.len(), 1);
+    assert_eq!(
+        validator_list.validators[0].status,
+        StakeStatus::ReadyForRemoval
+    );
+    assert_eq!(validator_list.validators[0].stake_lamports, 0);
+
+    let reserve_stake = context
+        .banks_client
+        .get_account(stake_pool_accounts.reserve_stake.pubkey())
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(
+        reserve_stake.lamports,
+        reserve_lamports + deactivated_lamports + stake_rent + 1
+    );
+
+    // Update stake pool balance, should be gone
+    let error = stake_pool_accounts
+        .update_stake_pool_balance(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+        )
+        .await;
+    assert!(error.is_none());
+
+    let validator_list = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let validator_list =
+        try_from_slice_unchecked::<ValidatorList>(validator_list.data.as_slice()).unwrap();
+    assert_eq!(validator_list.validators.len(), 0);
+}
 
 #[tokio::test]
 async fn fail_with_uninitialized_validator_list() {} // TODO

stake-pool/program/tests/vsa_add.rs

@@ -86,6 +86,7 @@ async fn success() {
             account_type: state::AccountType::ValidatorList,
             max_validators: stake_pool_accounts.max_validators,
             validators: vec![state::ValidatorStakeInfo {
+                status: state::StakeStatus::Active,
                 vote_account_address: user_stake.vote.pubkey(),
                 last_update_epoch: 0,
                 stake_lamports: 0,

stake-pool/program/tests/vsa_remove.rs

@@ -38,8 +38,8 @@ async fn setup() -> (
         .await
         .unwrap();
 
-    let user_stake = ValidatorStakeAccount::new(&stake_pool_accounts.stake_pool.pubkey());
+    let validator_stake = ValidatorStakeAccount::new(&stake_pool_accounts.stake_pool.pubkey());
-    user_stake
+    validator_stake
         .create_and_delegate(
             &mut banks_client,
             &payer,
@@ -53,7 +53,7 @@ async fn setup() -> (
             &mut banks_client,
             &payer,
             &recent_blockhash,
-            &user_stake.stake_account,
+            &validator_stake.stake_account,
         )
         .await;
     assert!(error.is_none());
@@ -63,13 +63,13 @@ async fn setup() -> (
         payer,
         recent_blockhash,
         stake_pool_accounts,
-        user_stake,
+        validator_stake,
     )
 }
 
 #[tokio::test]
 async fn success() {
-    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, user_stake) =
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake) =
         setup().await;
 
     let new_authority = Pubkey::new_unique();
@@ -79,8 +79,8 @@ async fn success() {
             &payer,
             &recent_blockhash,
             &new_authority,
-            &user_stake.stake_account,
+            &validator_stake.stake_account,
-            &user_stake.transient_stake_account,
+            &validator_stake.transient_stake_account,
         )
         .await;
     assert!(error.is_none());
@@ -103,7 +103,7 @@ async fn success() {
     );
 
     // Check of stake account authority has changed
-    let stake = get_account(&mut banks_client, &user_stake.stake_account).await;
+    let stake = get_account(&mut banks_client, &validator_stake.stake_account).await;
     let stake_state = deserialize::<stake_program::StakeState>(&stake.data).unwrap();
     match stake_state {
         stake_program::StakeState::Stake(meta, _) => {
@@ -116,7 +116,7 @@ async fn success() {
 
 #[tokio::test]
 async fn fail_with_wrong_stake_program_id() {
-    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, user_stake) =
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake) =
         setup().await;
 
     let wrong_stake_program = Pubkey::new_unique();
@@ -128,8 +128,8 @@ async fn fail_with_wrong_stake_program_id() {
         AccountMeta::new_readonly(stake_pool_accounts.withdraw_authority, false),
         AccountMeta::new_readonly(new_authority, false),
         AccountMeta::new(stake_pool_accounts.validator_list.pubkey(), false),
-        AccountMeta::new(user_stake.stake_account, false),
+        AccountMeta::new(validator_stake.stake_account, false),
-        AccountMeta::new_readonly(user_stake.transient_stake_account, false),
+        AccountMeta::new_readonly(validator_stake.transient_stake_account, false),
         AccountMeta::new_readonly(sysvar::clock::id(), false),
         AccountMeta::new_readonly(wrong_stake_program, false),
     ];
@@ -162,7 +162,7 @@ async fn fail_with_wrong_stake_program_id() {
 
 #[tokio::test]
 async fn fail_with_wrong_validator_list_account() {
-    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, user_stake) =
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake) =
         setup().await;
 
     let wrong_validator_list = Keypair::new();
@@ -176,8 +176,8 @@ async fn fail_with_wrong_validator_list_account() {
             &stake_pool_accounts.withdraw_authority,
             &new_authority,
             &wrong_validator_list.pubkey(),
-            &user_stake.stake_account,
+            &validator_stake.stake_account,
-            &user_stake.transient_stake_account,
+            &validator_stake.transient_stake_account,
         )
         .unwrap()],
         Some(&payer.pubkey()),
@@ -203,14 +203,14 @@ async fn fail_with_wrong_validator_list_account() {
 
 #[tokio::test]
 async fn fail_not_at_minimum() {
-    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, user_stake) =
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake) =
         setup().await;
 
     transfer(
         &mut banks_client,
         &payer,
         &recent_blockhash,
-        &user_stake.stake_account,
+        &validator_stake.stake_account,
         1_000_001,
     )
     .await;
@@ -222,8 +222,8 @@ async fn fail_not_at_minimum() {
             &payer,
             &recent_blockhash,
             &new_authority,
-            &user_stake.stake_account,
+            &validator_stake.stake_account,
-            &user_stake.transient_stake_account,
+            &validator_stake.transient_stake_account,
         )
         .await
         .unwrap()
@@ -239,7 +239,7 @@ async fn fail_not_at_minimum() {
 
 #[tokio::test]
 async fn fail_double_remove() {
-    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, user_stake) =
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake) =
         setup().await;
 
     let new_authority = Pubkey::new_unique();
@@ -249,8 +249,8 @@ async fn fail_double_remove() {
             &payer,
             &recent_blockhash,
             &new_authority,
-            &user_stake.stake_account,
+            &validator_stake.stake_account,
-            &user_stake.transient_stake_account,
+            &validator_stake.transient_stake_account,
         )
         .await;
     assert!(error.is_none());
@@ -263,8 +263,8 @@ async fn fail_double_remove() {
             &payer,
             &latest_blockhash,
             &new_authority,
-            &user_stake.stake_account,
+            &validator_stake.stake_account,
-            &user_stake.transient_stake_account,
+            &validator_stake.transient_stake_account,
         )
         .await
         .unwrap();
@@ -285,7 +285,7 @@ async fn fail_double_remove() {
 
 #[tokio::test]
 async fn fail_wrong_staker() {
-    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, user_stake) =
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake) =
         setup().await;
 
     let malicious = Keypair::new();
@@ -299,8 +299,8 @@ async fn fail_wrong_staker() {
             &stake_pool_accounts.withdraw_authority,
             &new_authority,
             &stake_pool_accounts.validator_list.pubkey(),
-            &user_stake.stake_account,
+            &validator_stake.stake_account,
-            &user_stake.transient_stake_account,
+            &validator_stake.transient_stake_account,
         )
         .unwrap()],
         Some(&payer.pubkey()),
@@ -328,7 +328,7 @@ async fn fail_wrong_staker() {
 
 #[tokio::test]
 async fn fail_no_signature() {
-    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, user_stake) =
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake) =
         setup().await;
 
     let new_authority = Pubkey::new_unique();
@@ -339,8 +339,8 @@ async fn fail_no_signature() {
         AccountMeta::new_readonly(stake_pool_accounts.withdraw_authority, false),
         AccountMeta::new_readonly(new_authority, false),
         AccountMeta::new(stake_pool_accounts.validator_list.pubkey(), false),
-        AccountMeta::new(user_stake.stake_account, false),
+        AccountMeta::new(validator_stake.stake_account, false),
-        AccountMeta::new_readonly(user_stake.transient_stake_account, false),
+        AccountMeta::new_readonly(validator_stake.transient_stake_account, false),
         AccountMeta::new_readonly(sysvar::clock::id(), false),
         AccountMeta::new_readonly(stake_program::id(), false),
     ];
@@ -378,7 +378,7 @@ async fn fail_no_signature() {
 
 #[tokio::test]
 async fn fail_with_activating_transient_stake() {
-    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, user_stake) =
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake) =
         setup().await;
 
     // increase the validator stake
@@ -387,8 +387,8 @@ async fn fail_with_activating_transient_stake() {
             &mut banks_client,
             &payer,
             &recent_blockhash,
-            &user_stake.transient_stake_account,
+            &validator_stake.transient_stake_account,
-            &user_stake.vote.pubkey(),
+            &validator_stake.vote.pubkey(),
             2_000_000_000,
         )
         .await;
@@ -401,8 +401,8 @@ async fn fail_with_activating_transient_stake() {
             &payer,
             &recent_blockhash,
             &new_authority,
-            &user_stake.stake_account,
+            &validator_stake.stake_account,
-            &user_stake.transient_stake_account,
+            &validator_stake.transient_stake_account,
         )
         .await
         .unwrap()
@@ -419,6 +419,129 @@ async fn fail_with_activating_transient_stake() {
     }
 }
 
+#[tokio::test]
+async fn success_with_deactivating_transient_stake() {
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake) =
+        setup().await;
+
+    let rent = banks_client.get_rent().await.unwrap();
+    let stake_rent = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
+    let deposit_info = simple_deposit(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts,
+        &validator_stake,
+        TEST_STAKE_AMOUNT,
+    )
+    .await
+    .unwrap();
+
+    // increase the validator stake
+    let error = stake_pool_accounts
+        .decrease_validator_stake(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &validator_stake.stake_account,
+            &validator_stake.transient_stake_account,
+            TEST_STAKE_AMOUNT + stake_rent,
+        )
+        .await;
+    assert!(error.is_none());
+
+    let new_authority = Pubkey::new_unique();
+    let error = stake_pool_accounts
+        .remove_validator_from_pool(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &new_authority,
+            &validator_stake.stake_account,
+            &validator_stake.transient_stake_account,
+        )
+        .await;
+    assert!(error.is_none());
+
+    // fail deposit
+    let maybe_deposit = simple_deposit(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts,
+        &validator_stake,
+        TEST_STAKE_AMOUNT,
+    )
+    .await;
+    assert!(maybe_deposit.is_none());
+
+    // fail withdraw
+    let user_stake_recipient = Keypair::new();
+    create_blank_stake_account(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &user_stake_recipient,
+    )
+    .await;
+
+    let new_authority = Pubkey::new_unique();
+    let error = stake_pool_accounts
+        .withdraw_stake(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &user_stake_recipient.pubkey(),
+            &deposit_info.pool_account.pubkey(),
+            &validator_stake.stake_account,
+            &new_authority,
+            1,
+        )
+        .await;
+    assert!(error.is_some());
+
+    // check validator has changed
+    let validator_list = get_account(
+        &mut banks_client,
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let validator_list =
+        try_from_slice_unchecked::<state::ValidatorList>(validator_list.data.as_slice()).unwrap();
+    let expected_list = state::ValidatorList {
+        account_type: state::AccountType::ValidatorList,
+        max_validators: stake_pool_accounts.max_validators,
+        validators: vec![state::ValidatorStakeInfo {
+            status: state::StakeStatus::DeactivatingTransient,
+            vote_account_address: validator_stake.vote.pubkey(),
+            last_update_epoch: 0,
+            stake_lamports: TEST_STAKE_AMOUNT + stake_rent,
+        }],
+    };
+    assert_eq!(validator_list, expected_list);
+
+    // Update, should not change, no merges yet
+    let error = stake_pool_accounts
+        .update_all(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &[validator_stake.vote.pubkey()],
+            false,
+        )
+        .await;
+    assert!(error.is_none());
+
+    let validator_list = get_account(
+        &mut banks_client,
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let validator_list =
+        try_from_slice_unchecked::<state::ValidatorList>(validator_list.data.as_slice()).unwrap();
+    assert_eq!(validator_list, expected_list);
+}
+
 #[tokio::test]
 async fn fail_not_updated_stake_pool() {} // TODO
 

stake-pool/program/tests/withdraw.rs

@@ -57,7 +57,8 @@ async fn setup() -> (
         &validator_stake_account,
         TEST_STAKE_AMOUNT,
     )
-    .await;
+    .await
+    .unwrap();
 
     let tokens_to_burn = deposit_info.pool_tokens / 4;
 
@@ -607,7 +608,8 @@ async fn fail_without_token_approval() {
         &validator_stake_account,
         TEST_STAKE_AMOUNT,
     )
-    .await;
+    .await
+    .unwrap();
 
     let tokens_to_burn = deposit_info.pool_tokens / 4;
 
@@ -675,7 +677,8 @@ async fn fail_with_low_delegation() {
         &validator_stake_account,
         TEST_STAKE_AMOUNT,
     )
-    .await;
+    .await
+    .unwrap();
 
     let tokens_to_burn = deposit_info.pool_tokens / 4;
 
@@ -819,7 +822,8 @@ async fn success_with_reserve() {
         &validator_stake,
         deposit_lamports,
     )
-    .await;
+    .await
+    .unwrap();
 
     // decrease some stake
     let error = stake_pool_accounts