diff --git a/programs/stake_api/src/stake_state.rs b/programs/stake_api/src/stake_state.rs index 41c7153c32..567c0e312c 100644 --- a/programs/stake_api/src/stake_state.rs +++ b/programs/stake_api/src/stake_state.rs @@ -563,6 +563,7 @@ impl<'a> StakeAccount for KeyedAccount<'a> { vote_account.account.lamports += voters_reward; stake.credits_observed = credits_observed; + stake.stake += stakers_reward; self.set_state(&StakeState::Stake(authorized, lockup, stake)) } else { @@ -1702,6 +1703,9 @@ mod tests { staker_rewards / 3 > voter_commission, "rewards should be split ~3:1" ); + // verify rewards are added to stake + let stake = StakeState::stake_from(&stake_keyed_account.account).unwrap(); + assert_eq!(stake.stake, stake_keyed_account.account.lamports); let wrong_vote_pubkey = Pubkey::new_rand(); let mut wrong_vote_keyed_account = diff --git a/programs/stake_tests/tests/stake_instruction.rs b/programs/stake_tests/tests/stake_instruction.rs index 7ba23b5af2..f0b6b3f6f8 100644 --- a/programs/stake_tests/tests/stake_instruction.rs +++ b/programs/stake_tests/tests/stake_instruction.rs @@ -5,18 +5,18 @@ use solana_runtime::{ genesis_utils::{create_genesis_block_with_leader, GenesisBlockInfo}, }; use solana_sdk::{ + account::Account, account_utils::State, client::SyncClient, message::Message, pubkey::Pubkey, signature::{Keypair, KeypairUtil}, - sysvar, - sysvar::rewards::Rewards, + sysvar::{self, rewards::Rewards}, }; use solana_stake_api::{ id, stake_instruction::{self, process_instruction}, - stake_state::{self, StakeState}, + stake_state::{self, Stake, StakeState}, }; use solana_vote_api::{ vote_instruction, @@ -101,7 +101,7 @@ fn test_stake_account_delegate() { &mint_pubkey, &staker_pubkey, &vote_pubkey, - 20000, + 1_000_000, &authorized, )); bank_client @@ -112,7 +112,7 @@ fn test_stake_account_delegate() { let account = bank.get_account(&staker_pubkey).expect("account not found"); let stake_state = account.state().expect("couldn't unpack account data"); if let StakeState::Stake(_authorized, _lockup, stake) = stake_state { - assert_eq!(stake.stake, 20000); + assert_eq!(stake.stake, 1_000_000); } else { assert!(false, "wrong account type found") } @@ -122,7 +122,7 @@ fn test_stake_account_delegate() { vec![stake_instruction::withdraw( &staker_pubkey, &Pubkey::new_rand(), - 20000, + 1_000_000, )], Some(&mint_pubkey), ); @@ -134,7 +134,7 @@ fn test_stake_account_delegate() { let account = bank.get_account(&staker_pubkey).expect("account not found"); let stake_state = account.state().expect("couldn't unpack account data"); if let StakeState::Stake(_authorized, _lockup, stake) = stake_state { - assert_eq!(stake.stake, 20000); + assert_eq!(stake.stake, 1_000_000); } else { assert!(false, "wrong account type found") } @@ -173,19 +173,17 @@ fn test_stake_account_delegate() { ); assert_matches!(bank_client.send_message(&[&mint_keypair], message), Ok(_)); - // Test that balance increased, and calculate the rewards - let rewards; - let account = bank.get_account(&staker_pubkey).expect("account not found"); - let stake_state = account.state().expect("couldn't unpack account data"); - if let StakeState::Stake(_authorized, _lockup, stake) = stake_state { - assert!(account.lamports > 20000); - assert_eq!(stake.stake, 20000); - rewards = account.lamports - 20000; - } else { - rewards = 0; - assert!(false, "wrong account type found") + fn get_stake(bank: &Bank, staker_pubkey: &Pubkey) -> (Stake, Account) { + let account = bank.get_account(staker_pubkey).unwrap(); + (StakeState::stake_from(&account).unwrap(), account) } + // Test that balance increased, and calculate the rewards + let (stake, account) = get_stake(&bank, &staker_pubkey); + assert!(account.lamports > 1_000_000); + assert!(stake.stake > 1_000_000); + let rewards = account.lamports - 1_000_000; + // Deactivate the stake let message = Message::new_with_payer( vec![stake_instruction::deactivate_stake( @@ -198,47 +196,59 @@ fn test_stake_account_delegate() { .send_message(&[&mint_keypair, &staker_keypair], message) .is_ok()); - // // Test that we cannot withdraw staked lamports due to cooldown period - // let message = Message::new_with_payer( - // vec![stake_instruction::withdraw( - // &staker_pubkey, - // &Pubkey::new_rand(), - // 20000, - // )], - // Some(&mint_pubkey), - // ); - // assert!(bank_client - // .send_message(&[&mint_keypair, &staker_keypair], message) - // .is_err()); - // - // let old_epoch = bank.epoch(); - // let slots = bank.get_slots_in_epoch(old_epoch); - // - // // Create a new bank at later epoch (within cooldown period) - // let bank = Bank::new_from_parent(&bank, &Pubkey::default(), slots + bank.slot()); - // assert_ne!(old_epoch, bank.epoch()); - // let bank = Arc::new(bank); - // let bank_client = BankClient::new_shared(&bank); - // - // let message = Message::new_with_payer( - // vec![stake_instruction::withdraw( - // &staker_pubkey, - // &Pubkey::new_rand(), - // 20000, - // )], - // Some(&mint_pubkey), - // ); - // assert!(bank_client - // .send_message(&[&mint_keypair, &staker_keypair], message) - // .is_err()); - // TODO: implement cooldown + // Test that we cannot withdraw staked lamports due to cooldown period + let message = Message::new_with_payer( + vec![stake_instruction::withdraw( + &staker_pubkey, + &Pubkey::new_rand(), + 1_000_000, + )], + Some(&mint_pubkey), + ); + assert!(bank_client + .send_message(&[&mint_keypair, &staker_keypair], message) + .is_err()); - // Create a new bank at later epoch (to account for cooldown of stake) + let old_epoch = bank.epoch(); + let slots = bank.get_slots_in_epoch(old_epoch); + // Create a new bank at later epoch (within cooldown period) + let bank = Bank::new_from_parent(&bank, &Pubkey::default(), slots + bank.slot()); + assert_ne!(old_epoch, bank.epoch()); + let bank = Arc::new(bank); + let bank_client = BankClient::new_shared(&bank); + + let message = Message::new_with_payer( + vec![stake_instruction::withdraw( + &staker_pubkey, + &Pubkey::new_rand(), + 1_000_000, + )], + Some(&mint_pubkey), + ); + + assert!(bank_client + .send_message(&[&mint_keypair, &staker_keypair], message) + .is_err()); + + let message = Message::new_with_payer( + vec![stake_instruction::withdraw( + &staker_pubkey, + &Pubkey::new_rand(), + 250_000, + )], + Some(&mint_pubkey), + ); + // assert we can withdraw some smaller amount, more than rewards + assert!(bank_client + .send_message(&[&mint_keypair, &staker_keypair], message) + .is_ok()); + + // Create a new bank at a much later epoch (to account for cooldown of stake) let mut bank = Bank::new_from_parent( &bank, &Pubkey::default(), - genesis_block.slots_per_epoch * 4 + bank.slot(), + genesis_block.slots_per_epoch * 10 + bank.slot(), ); bank.add_instruction_processor(id(), process_instruction); let bank = Arc::new(bank); @@ -249,7 +259,7 @@ fn test_stake_account_delegate() { vec![stake_instruction::withdraw( &staker_pubkey, &Pubkey::new_rand(), - 20000, + 750_000, )], Some(&mint_pubkey), ); @@ -258,11 +268,6 @@ fn test_stake_account_delegate() { .is_ok()); // Test that balance and stake is updated correctly (we have withdrawn all lamports except rewards) - let account = bank.get_account(&staker_pubkey).expect("account not found"); - let stake_state = account.state().expect("couldn't unpack account data"); - if let StakeState::Stake(_, _, _stake) = stake_state { - assert_eq!(account.lamports, rewards); - } else { - assert!(false, "wrong account type found") - } + let (_stake, account) = get_stake(&bank, &staker_pubkey); + assert_eq!(account.lamports, rewards); }