|
2 | 2 | super::{ |
3 | 3 | instruction_builders::{InstructionConfig, InstructionExecution}, |
4 | 4 | lifecycle::StakeLifecycle, |
5 | | - utils::{add_sysvars, STAKE_RENT_EXEMPTION}, |
| 5 | + utils::{add_sysvars, create_vote_account, STAKE_RENT_EXEMPTION}, |
6 | 6 | }, |
7 | 7 | mollusk_svm::{result::Check, Mollusk}, |
8 | 8 | solana_account::AccountSharedData, |
9 | 9 | solana_instruction::Instruction, |
10 | 10 | solana_pubkey::Pubkey, |
| 11 | + solana_stake_interface::state::Lockup, |
11 | 12 | solana_stake_program::id, |
12 | 13 | }; |
13 | 14 |
|
14 | 15 | /// Builder for creating stake accounts with customizable parameters |
15 | | -pub struct StakeAccountBuilder { |
| 16 | +pub struct StakeAccountBuilder<'a> { |
| 17 | + ctx: &'a mut StakeTestContext, |
16 | 18 | lifecycle: StakeLifecycle, |
| 19 | + staked_amount: u64, |
| 20 | + stake_authority: Option<Pubkey>, |
| 21 | + withdraw_authority: Option<Pubkey>, |
| 22 | + lockup: Option<Lockup>, |
| 23 | + vote_account: Option<Pubkey>, |
| 24 | + stake_pubkey: Option<Pubkey>, |
17 | 25 | } |
18 | 26 |
|
19 | | -impl StakeAccountBuilder { |
| 27 | +impl StakeAccountBuilder<'_> { |
| 28 | + /// Set the staked amount (lamports delegated to validator) |
| 29 | + pub fn staked_amount(mut self, amount: u64) -> Self { |
| 30 | + self.staked_amount = amount; |
| 31 | + self |
| 32 | + } |
| 33 | + |
| 34 | + /// Set a custom stake authority (defaults to ctx.staker) |
| 35 | + pub fn stake_authority(mut self, authority: &Pubkey) -> Self { |
| 36 | + self.stake_authority = Some(*authority); |
| 37 | + self |
| 38 | + } |
| 39 | + |
| 40 | + /// Set a custom withdraw authority (defaults to ctx.withdrawer) |
| 41 | + pub fn withdraw_authority(mut self, authority: &Pubkey) -> Self { |
| 42 | + self.withdraw_authority = Some(*authority); |
| 43 | + self |
| 44 | + } |
| 45 | + |
| 46 | + /// Set a custom lockup (defaults to Lockup::default()) |
| 47 | + pub fn lockup(mut self, lockup: &Lockup) -> Self { |
| 48 | + self.lockup = Some(*lockup); |
| 49 | + self |
| 50 | + } |
| 51 | + |
| 52 | + /// Set a custom vote account (defaults to ctx.vote_account) |
| 53 | + pub fn vote_account(mut self, vote_account: &Pubkey) -> Self { |
| 54 | + self.vote_account = Some(*vote_account); |
| 55 | + self |
| 56 | + } |
| 57 | + |
| 58 | + /// Set a specific stake account pubkey (defaults to Pubkey::new_unique()) |
| 59 | + pub fn stake_pubkey(mut self, pubkey: &Pubkey) -> Self { |
| 60 | + self.stake_pubkey = Some(*pubkey); |
| 61 | + self |
| 62 | + } |
| 63 | + |
| 64 | + /// Build the stake account and return (pubkey, account_data) |
20 | 65 | pub fn build(self) -> (Pubkey, AccountSharedData) { |
21 | | - let stake_pubkey = Pubkey::new_unique(); |
22 | | - let account = self.lifecycle.create_uninitialized_account(); |
| 66 | + let stake_pubkey = self.stake_pubkey.unwrap_or_else(Pubkey::new_unique); |
| 67 | + let account = self.lifecycle.create_stake_account_fully_specified( |
| 68 | + &mut self.ctx.mollusk, |
| 69 | + &stake_pubkey, |
| 70 | + self.vote_account.as_ref().unwrap_or( |
| 71 | + self.ctx |
| 72 | + .vote_account |
| 73 | + .as_ref() |
| 74 | + .expect("vote_account required for this lifecycle"), |
| 75 | + ), |
| 76 | + self.staked_amount, |
| 77 | + self.stake_authority.as_ref().unwrap_or(&self.ctx.staker), |
| 78 | + self.withdraw_authority |
| 79 | + .as_ref() |
| 80 | + .unwrap_or(&self.ctx.withdrawer), |
| 81 | + self.lockup.as_ref().unwrap_or(&Lockup::default()), |
| 82 | + ); |
23 | 83 | (stake_pubkey, account) |
24 | 84 | } |
25 | 85 | } |
26 | 86 |
|
27 | | -/// Consolidated test context for stake account tests |
28 | 87 | pub struct StakeTestContext { |
29 | 88 | pub mollusk: Mollusk, |
30 | 89 | pub rent_exempt_reserve: u64, |
31 | 90 | pub staker: Pubkey, |
32 | 91 | pub withdrawer: Pubkey, |
| 92 | + pub minimum_delegation: Option<u64>, |
| 93 | + pub vote_account: Option<Pubkey>, |
| 94 | + pub vote_account_data: Option<AccountSharedData>, |
33 | 95 | } |
34 | 96 |
|
35 | 97 | impl StakeTestContext { |
36 | | - /// Create a new test context with all standard setup |
37 | | - pub fn new() -> Self { |
| 98 | + pub fn minimal() -> Self { |
38 | 99 | let mollusk = Mollusk::new(&id(), "solana_stake_program"); |
| 100 | + Self { |
| 101 | + mollusk, |
| 102 | + rent_exempt_reserve: STAKE_RENT_EXEMPTION, |
| 103 | + staker: Pubkey::new_unique(), |
| 104 | + withdrawer: Pubkey::new_unique(), |
| 105 | + minimum_delegation: None, |
| 106 | + vote_account: None, |
| 107 | + vote_account_data: None, |
| 108 | + } |
| 109 | + } |
39 | 110 |
|
| 111 | + pub fn with_delegation() -> Self { |
| 112 | + let mollusk = Mollusk::new(&id(), "solana_stake_program"); |
| 113 | + let minimum_delegation = solana_stake_program::get_minimum_delegation(); |
40 | 114 | Self { |
41 | 115 | mollusk, |
42 | 116 | rent_exempt_reserve: STAKE_RENT_EXEMPTION, |
43 | 117 | staker: Pubkey::new_unique(), |
44 | 118 | withdrawer: Pubkey::new_unique(), |
| 119 | + minimum_delegation: Some(minimum_delegation), |
| 120 | + vote_account: Some(Pubkey::new_unique()), |
| 121 | + vote_account_data: Some(create_vote_account()), |
45 | 122 | } |
46 | 123 | } |
47 | 124 |
|
| 125 | + pub fn new() -> Self { |
| 126 | + Self::with_delegation() |
| 127 | + } |
| 128 | + |
48 | 129 | /// Create a stake account builder for the specified lifecycle stage |
49 | 130 | /// |
50 | 131 | /// Example: |
51 | 132 | /// ``` |
52 | 133 | /// let (stake, account) = ctx |
53 | | - /// .stake_account(StakeLifecycle::Uninitialized) |
| 134 | + /// .stake_account(StakeLifecycle::Active) |
| 135 | + /// .staked_amount(1_000_000) |
54 | 136 | /// .build(); |
55 | 137 | /// ``` |
56 | 138 | pub fn stake_account(&mut self, lifecycle: StakeLifecycle) -> StakeAccountBuilder { |
57 | | - StakeAccountBuilder { lifecycle } |
| 139 | + StakeAccountBuilder { |
| 140 | + ctx: self, |
| 141 | + lifecycle, |
| 142 | + staked_amount: 0, |
| 143 | + stake_authority: None, |
| 144 | + withdraw_authority: None, |
| 145 | + lockup: None, |
| 146 | + vote_account: None, |
| 147 | + stake_pubkey: None, |
| 148 | + } |
| 149 | + } |
| 150 | + |
| 151 | + /// Create a lockup that expires in the future |
| 152 | + pub fn create_future_lockup(&self, epochs_ahead: u64) -> Lockup { |
| 153 | + Lockup { |
| 154 | + unix_timestamp: 0, |
| 155 | + epoch: self.mollusk.sysvars.clock.epoch + epochs_ahead, |
| 156 | + custodian: Pubkey::new_unique(), |
| 157 | + } |
| 158 | + } |
| 159 | + |
| 160 | + /// Create a lockup that's currently in force (far future) |
| 161 | + pub fn create_in_force_lockup(&self) -> Lockup { |
| 162 | + self.create_future_lockup(1_000_000) |
| 163 | + } |
| 164 | + |
| 165 | + /// Create a second vote account (for testing different vote accounts) |
| 166 | + pub fn create_second_vote_account(&self) -> (Pubkey, AccountSharedData) { |
| 167 | + (Pubkey::new_unique(), create_vote_account()) |
58 | 168 | } |
59 | 169 |
|
60 | | - /// Process an instruction |
| 170 | + /// Process an instruction with a config-based approach |
61 | 171 | pub fn process_with<'b, C: InstructionConfig>( |
62 | 172 | &self, |
63 | 173 | config: C, |
|
0 commit comments