From 423127d91eabef1b49bc9d059a57b9b90d7c858b Mon Sep 17 00:00:00 2001 From: rustopian <96253492+rustopian@users.noreply.github.com> Date: Wed, 17 Dec 2025 17:42:26 +0000 Subject: [PATCH 1/4] basic helpers and Initialize(Checked) tests --- Cargo.lock | 1 + program/Cargo.toml | 1 + program/tests/helpers/context.rs | 109 +++++ program/tests/helpers/instruction_builders.rs | 61 +++ program/tests/helpers/lifecycle.rs | 29 ++ program/tests/helpers/mod.rs | 7 + program/tests/helpers/utils.rs | 65 +++ program/tests/initialize.rs | 410 ++++++++++++++++++ program/tests/program_test.rs | 121 ------ 9 files changed, 683 insertions(+), 121 deletions(-) create mode 100644 program/tests/helpers/context.rs create mode 100644 program/tests/helpers/instruction_builders.rs create mode 100644 program/tests/helpers/lifecycle.rs create mode 100644 program/tests/helpers/mod.rs create mode 100644 program/tests/helpers/utils.rs create mode 100644 program/tests/initialize.rs diff --git a/Cargo.lock b/Cargo.lock index d3ec163c..aa7f5eed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6684,6 +6684,7 @@ dependencies = [ "solana-sdk-ids", "solana-signature", "solana-signer", + "solana-stake-client", "solana-stake-interface 2.0.1", "solana-svm-log-collector", "solana-system-interface", diff --git a/program/Cargo.toml b/program/Cargo.toml index 97c5ae46..21d326e9 100644 --- a/program/Cargo.toml +++ b/program/Cargo.toml @@ -43,6 +43,7 @@ solana-sdk-ids = "3.0.0" solana-signature = "3.0.0" solana-signer = "3.0.0" solana-svm-log-collector = "3.0.0" +solana-stake-client = { path = "../clients/rust" } solana-system-interface = { version = "2.0.0", features = ["bincode"] } solana-transaction = "3.0.0" test-case = "3.3.1" diff --git a/program/tests/helpers/context.rs b/program/tests/helpers/context.rs new file mode 100644 index 00000000..ccd3faf5 --- /dev/null +++ b/program/tests/helpers/context.rs @@ -0,0 +1,109 @@ +use { + super::{ + instruction_builders::InstructionExecution, + lifecycle::StakeLifecycle, + utils::{add_sysvars, STAKE_RENT_EXEMPTION}, + }, + mollusk_svm::{result::Check, Mollusk}, + solana_account::AccountSharedData, + solana_instruction::Instruction, + solana_pubkey::Pubkey, + solana_stake_program::id, +}; + +/// Builder for creating stake accounts with customizable parameters +pub struct StakeAccountBuilder { + lifecycle: StakeLifecycle, +} + +impl StakeAccountBuilder { + pub fn build(self) -> (Pubkey, AccountSharedData) { + let stake_pubkey = Pubkey::new_unique(); + let account = self.lifecycle.create_uninitialized_account(); + (stake_pubkey, account) + } +} + +/// Consolidated test context for stake account tests +pub struct StakeTestContext { + pub mollusk: Mollusk, + pub rent_exempt_reserve: u64, + pub staker: Pubkey, + pub withdrawer: Pubkey, +} + +impl StakeTestContext { + /// Create a new test context with all standard setup + pub fn new() -> Self { + let mollusk = Mollusk::new(&id(), "solana_stake_program"); + + Self { + mollusk, + rent_exempt_reserve: STAKE_RENT_EXEMPTION, + staker: Pubkey::new_unique(), + withdrawer: Pubkey::new_unique(), + } + } + + /// Create a stake account builder for the specified lifecycle stage + /// + /// Example: + /// ``` + /// let (stake, account) = ctx + /// .stake_account(StakeLifecycle::Uninitialized) + /// .build(); + /// ``` + pub fn stake_account(&mut self, lifecycle: StakeLifecycle) -> StakeAccountBuilder { + StakeAccountBuilder { lifecycle } + } + + /// Process an instruction + pub fn process_with( + &self, + instruction: Instruction, + accounts: Vec<(Pubkey, AccountSharedData)>, + ) -> InstructionExecution { + InstructionExecution::new(instruction, accounts, self) + } + + /// Process an instruction with optional missing signer testing + pub(crate) fn process_instruction_maybe_test_signers( + &self, + instruction: &Instruction, + accounts: Vec<(Pubkey, AccountSharedData)>, + checks: &[Check], + test_missing_signers: bool, + ) -> mollusk_svm::result::InstructionResult { + if test_missing_signers { + use solana_program_error::ProgramError; + + // Test that removing each signer causes failure + for i in 0..instruction.accounts.len() { + if instruction.accounts[i].is_signer { + let mut modified_instruction = instruction.clone(); + modified_instruction.accounts[i].is_signer = false; + + let accounts_with_sysvars = + add_sysvars(&self.mollusk, &modified_instruction, accounts.clone()); + + self.mollusk.process_and_validate_instruction( + &modified_instruction, + &accounts_with_sysvars, + &[Check::err(ProgramError::MissingRequiredSignature)], + ); + } + } + } + + // Process with all signers present + let accounts_with_sysvars = add_sysvars(&self.mollusk, instruction, accounts); + self.mollusk + .process_and_validate_instruction(instruction, &accounts_with_sysvars, checks) + } +} + +impl Default for StakeTestContext { + fn default() -> Self { + Self::new() + } +} diff --git a/program/tests/helpers/instruction_builders.rs b/program/tests/helpers/instruction_builders.rs new file mode 100644 index 00000000..f805749d --- /dev/null +++ b/program/tests/helpers/instruction_builders.rs @@ -0,0 +1,61 @@ +use { + super::context::StakeTestContext, mollusk_svm::result::Check, + solana_account::AccountSharedData, solana_instruction::Instruction, solana_pubkey::Pubkey, +}; + +/// Execution builder with validation and signer testing +pub struct InstructionExecution<'a, 'b> { + instruction: Instruction, + accounts: Vec<(Pubkey, AccountSharedData)>, + ctx: &'a StakeTestContext, + checks: Option<&'b [Check<'b>]>, + test_missing_signers: Option, // `None` runs if `Check::success` +} + +impl<'b> InstructionExecution<'_, 'b> { + pub fn checks(mut self, checks: &'b [Check<'b>]) -> Self { + self.checks = Some(checks); + self + } + + pub fn test_missing_signers(mut self, test: bool) -> Self { + self.test_missing_signers = Some(test); + self + } + + /// Executes the instruction. If `checks` is `None` or empty, uses `Check::success()`. + /// Fail-safe default: when `test_missing_signers` is `None`, runs the missing-signers + /// test (`true`). Callers must explicitly opt out with `.test_missing_signers(false)`. + pub fn execute(self) -> mollusk_svm::result::InstructionResult { + let default_checks = [Check::success()]; + let checks = match self.checks { + Some(c) if !c.is_empty() => c, + _ => &default_checks, + }; + + let test_missing_signers = self.test_missing_signers.unwrap_or(true); + + self.ctx.process_instruction_maybe_test_signers( + &self.instruction, + self.accounts, + checks, + test_missing_signers, + ) + } +} + +impl<'a> InstructionExecution<'a, '_> { + pub(crate) fn new( + instruction: Instruction, + accounts: Vec<(Pubkey, AccountSharedData)>, + ctx: &'a StakeTestContext, + ) -> Self { + Self { + instruction, + accounts, + ctx, + checks: None, + test_missing_signers: None, + } + } +} diff --git a/program/tests/helpers/lifecycle.rs b/program/tests/helpers/lifecycle.rs new file mode 100644 index 00000000..74e55285 --- /dev/null +++ b/program/tests/helpers/lifecycle.rs @@ -0,0 +1,29 @@ +use { + super::utils::STAKE_RENT_EXEMPTION, solana_account::AccountSharedData, + solana_stake_interface::state::StakeStateV2, solana_stake_program::id, +}; + +/// Lifecycle states for stake accounts in tests +#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub enum StakeLifecycle { + Uninitialized = 0, + Initialized, + Activating, + Active, + Deactivating, + Deactive, + Closed, +} + +impl StakeLifecycle { + /// Create an uninitialized stake account + pub fn create_uninitialized_account(self) -> AccountSharedData { + AccountSharedData::new_data_with_space( + STAKE_RENT_EXEMPTION, + &StakeStateV2::Uninitialized, + StakeStateV2::size_of(), + &id(), + ) + .unwrap() + } +} diff --git a/program/tests/helpers/mod.rs b/program/tests/helpers/mod.rs new file mode 100644 index 00000000..ff09b000 --- /dev/null +++ b/program/tests/helpers/mod.rs @@ -0,0 +1,7 @@ +#![allow(clippy::arithmetic_side_effects)] +#![allow(dead_code)] + +pub mod context; +pub mod instruction_builders; +pub mod lifecycle; +pub mod utils; diff --git a/program/tests/helpers/utils.rs b/program/tests/helpers/utils.rs new file mode 100644 index 00000000..3bc0659b --- /dev/null +++ b/program/tests/helpers/utils.rs @@ -0,0 +1,65 @@ +use { + mollusk_svm::Mollusk, + solana_account::{Account, AccountSharedData}, + solana_instruction::Instruction, + solana_pubkey::Pubkey, + solana_rent::Rent, + solana_stake_interface::{stake_history::StakeHistory, state::StakeStateV2}, + solana_sysvar_id::SysvarId, + std::collections::HashMap, +}; + +// hardcoded for convenience +pub const STAKE_RENT_EXEMPTION: u64 = 2_282_880; + +#[test] +fn assert_stake_rent_exemption() { + assert_eq!( + Rent::default().minimum_balance(StakeStateV2::size_of()), + STAKE_RENT_EXEMPTION + ); +} + +/// Resolve all accounts for an instruction, including sysvars and instruction accounts +/// +/// This function re-serializes the stake history sysvar from mollusk.sysvars.stake_history +/// every time it's called, ensuring that any updates to the stake history are reflected in the accounts. +pub fn add_sysvars( + mollusk: &Mollusk, + instruction: &Instruction, + accounts: Vec<(Pubkey, AccountSharedData)>, +) -> Vec<(Pubkey, Account)> { + // Build a map of provided accounts + let mut account_map: HashMap = accounts + .into_iter() + .map(|(pk, acc)| (pk, acc.into())) + .collect(); + + // Now resolve all accounts from the instruction + let mut result = Vec::new(); + for account_meta in &instruction.accounts { + let key = account_meta.pubkey; + let account = if let Some(acc) = account_map.remove(&key) { + // Use the provided account + acc + } else if Rent::check_id(&key) { + mollusk.sysvars.keyed_account_for_rent_sysvar().1 + } else if solana_clock::Clock::check_id(&key) { + mollusk.sysvars.keyed_account_for_clock_sysvar().1 + } else if solana_epoch_schedule::EpochSchedule::check_id(&key) { + mollusk.sysvars.keyed_account_for_epoch_schedule_sysvar().1 + } else if solana_epoch_rewards::EpochRewards::check_id(&key) { + mollusk.sysvars.keyed_account_for_epoch_rewards_sysvar().1 + } else if StakeHistory::check_id(&key) { + // Re-serialize stake history from mollusk.sysvars.stake_history + mollusk.sysvars.keyed_account_for_stake_history_sysvar().1 + } else { + // Default empty account + Account::default() + }; + + result.push((key, account)); + } + + result +} diff --git a/program/tests/initialize.rs b/program/tests/initialize.rs new file mode 100644 index 00000000..debc92b3 --- /dev/null +++ b/program/tests/initialize.rs @@ -0,0 +1,410 @@ +#![allow(clippy::arithmetic_side_effects)] + +mod helpers; + +use { + helpers::{ + context::StakeTestContext, instruction_builders::InstructionExecution, + lifecycle::StakeLifecycle, + }, + mollusk_svm::result::Check, + solana_account::{AccountSharedData, ReadableAccount}, + solana_program_error::ProgramError, + solana_pubkey::Pubkey, + solana_rent::Rent, + solana_sdk_ids::{stake::id, system_program::id as system_program_id}, + solana_stake_client::instructions::{InitializeBuilder, InitializeCheckedBuilder}, + solana_stake_interface::state::StakeStateV2, + test_case::test_case, +}; + +#[derive(Debug, Clone, Copy)] +enum InitializeVariant { + Initialize, + InitializeChecked, +} + +fn lockup_for( + variant: InitializeVariant, + custodian: Pubkey, +) -> solana_stake_interface::state::Lockup { + match variant { + InitializeVariant::Initialize => solana_stake_interface::state::Lockup { + epoch: 1, + unix_timestamp: 0, + custodian, + }, + InitializeVariant::InitializeChecked => solana_stake_interface::state::Lockup::default(), + } +} + +/// Extension trait for InitializeBuilder to add accounts helper +trait InitializeBuilderExt { + fn with_stake_account( + self, + stake: &Pubkey, + stake_account: &AccountSharedData, + ) -> (Self, Vec<(Pubkey, AccountSharedData)>) + where + Self: Sized; +} + +impl InitializeBuilderExt for InitializeBuilder { + fn with_stake_account( + mut self, + stake: &Pubkey, + stake_account: &AccountSharedData, + ) -> (Self, Vec<(Pubkey, AccountSharedData)>) { + self.stake(*stake); + let accounts = vec![(*stake, stake_account.clone())]; + (self, accounts) + } +} + +/// Extension trait for InitializeCheckedBuilder to add accounts helper +trait InitializeCheckedBuilderExt { + fn with_stake_account( + self, + stake: &Pubkey, + stake_account: &AccountSharedData, + ) -> (Self, Vec<(Pubkey, AccountSharedData)>) + where + Self: Sized; +} + +impl InitializeCheckedBuilderExt for InitializeCheckedBuilder { + fn with_stake_account( + mut self, + stake: &Pubkey, + stake_account: &AccountSharedData, + ) -> (Self, Vec<(Pubkey, AccountSharedData)>) { + self.stake(*stake); + let accounts = vec![(*stake, stake_account.clone())]; + (self, accounts) + } +} + +#[test_case(InitializeVariant::Initialize; "initialize")] +#[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] +fn test_initialize(variant: InitializeVariant) { + let mut ctx = StakeTestContext::new(); + + let staker = ctx.staker; + let withdrawer = ctx.withdrawer; + let lockup = lockup_for(variant, Pubkey::new_unique()); + + let (stake, stake_account) = ctx.stake_account(StakeLifecycle::Uninitialized).build(); + + let result = { + let program_id = id(); + let checks = [ + Check::success(), + Check::all_rent_exempt(), + Check::account(&stake) + .lamports(ctx.rent_exempt_reserve) + .owner(&program_id) + .space(StakeStateV2::size_of()) + .build(), + ]; + + let processor: InstructionExecution = match variant { + InitializeVariant::Initialize => { + let (mut builder, accounts) = + InitializeBuilder::new().with_stake_account(&stake, &stake_account); + + let instruction = builder + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(); + + ctx.process_with(instruction, accounts) + } + InitializeVariant::InitializeChecked => { + let (mut builder, accounts) = + InitializeCheckedBuilder::new().with_stake_account(&stake, &stake_account); + + let instruction = builder + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(); + + ctx.process_with(instruction, accounts) + } + }; + + processor + .checks(&checks) + .test_missing_signers(true) + .execute() + }; + + let resulting_account: AccountSharedData = result.resulting_accounts[0].1.clone().into(); + let stake_state: StakeStateV2 = bincode::deserialize(resulting_account.data()).unwrap(); + assert_eq!( + stake_state, + StakeStateV2::Initialized(solana_stake_interface::state::Meta { + authorized: solana_stake_interface::state::Authorized { staker, withdrawer }, + rent_exempt_reserve: ctx.rent_exempt_reserve, + lockup, + }), + ); + + // Re-initialize should fail + let processor: InstructionExecution = match variant { + InitializeVariant::Initialize => { + let (mut builder, accounts) = + InitializeBuilder::new().with_stake_account(&stake, &resulting_account); + + let instruction = builder + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(); + + ctx.process_with(instruction, accounts) + } + InitializeVariant::InitializeChecked => { + let (mut builder, accounts) = + InitializeCheckedBuilder::new().with_stake_account(&stake, &resulting_account); + + let instruction = builder + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(); + + ctx.process_with(instruction, accounts) + } + }; + + processor + .checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute(); +} + +#[test_case(InitializeVariant::Initialize; "initialize")] +#[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] +fn test_initialize_insufficient_funds(variant: InitializeVariant) { + let ctx = StakeTestContext::new(); + + let staker = ctx.staker; + let withdrawer = ctx.withdrawer; + let lockup = lockup_for(variant, Pubkey::new_unique()); + + // Account has insufficient lamports + let stake = Pubkey::new_unique(); + let stake_account = AccountSharedData::new_data_with_space( + ctx.rent_exempt_reserve / 2, + &StakeStateV2::Uninitialized, + StakeStateV2::size_of(), + &id(), + ) + .unwrap(); + + let processor: InstructionExecution = match variant { + InitializeVariant::Initialize => { + let (mut builder, accounts) = + InitializeBuilder::new().with_stake_account(&stake, &stake_account); + + let instruction = builder + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(); + + ctx.process_with(instruction, accounts) + } + InitializeVariant::InitializeChecked => { + let (mut builder, accounts) = + InitializeCheckedBuilder::new().with_stake_account(&stake, &stake_account); + + let instruction = builder + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(); + + ctx.process_with(instruction, accounts) + } + }; + + processor + .checks(&[Check::err(ProgramError::InsufficientFunds)]) + .test_missing_signers(false) + .execute(); +} + +#[test_case(InitializeVariant::Initialize; "initialize")] +#[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] +fn test_initialize_incorrect_size_larger(variant: InitializeVariant) { + let ctx = StakeTestContext::new(); + + let rent_exempt_reserve = Rent::default().minimum_balance(StakeStateV2::size_of() * 2); + + let staker = ctx.staker; + let withdrawer = ctx.withdrawer; + let lockup = lockup_for(variant, Pubkey::new_unique()); + + // Account data length too large + let stake = Pubkey::new_unique(); + let stake_account = AccountSharedData::new_data_with_space( + rent_exempt_reserve, + &StakeStateV2::Uninitialized, + StakeStateV2::size_of() + 1, + &id(), + ) + .unwrap(); + + let processor: InstructionExecution = match variant { + InitializeVariant::Initialize => { + let (mut builder, accounts) = + InitializeBuilder::new().with_stake_account(&stake, &stake_account); + + let instruction = builder + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(); + + ctx.process_with(instruction, accounts) + } + InitializeVariant::InitializeChecked => { + let (mut builder, accounts) = + InitializeCheckedBuilder::new().with_stake_account(&stake, &stake_account); + + let instruction = builder + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(); + + ctx.process_with(instruction, accounts) + } + }; + + processor + .checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute(); +} + +#[test_case(InitializeVariant::Initialize; "initialize")] +#[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] +fn test_initialize_incorrect_size_smaller(variant: InitializeVariant) { + let ctx = StakeTestContext::new(); + + let rent_exempt_reserve = Rent::default().minimum_balance(StakeStateV2::size_of()); + + let staker = ctx.staker; + let withdrawer = ctx.withdrawer; + let lockup = lockup_for(variant, Pubkey::new_unique()); + + // Account data length too small + let stake = Pubkey::new_unique(); + let stake_account = AccountSharedData::new_data_with_space( + rent_exempt_reserve, + &StakeStateV2::Uninitialized, + StakeStateV2::size_of() - 1, + &id(), + ) + .unwrap(); + + let processor: InstructionExecution = match variant { + InitializeVariant::Initialize => { + let (mut builder, accounts) = + InitializeBuilder::new().with_stake_account(&stake, &stake_account); + + let instruction = builder + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(); + + ctx.process_with(instruction, accounts) + } + InitializeVariant::InitializeChecked => { + let (mut builder, accounts) = + InitializeCheckedBuilder::new().with_stake_account(&stake, &stake_account); + + let instruction = builder + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(); + + ctx.process_with(instruction, accounts) + } + }; + + processor + .checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute(); +} + +#[test_case(InitializeVariant::Initialize; "initialize")] +#[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] +fn test_initialize_wrong_owner(variant: InitializeVariant) { + let ctx = StakeTestContext::new(); + + let staker = ctx.staker; + let withdrawer = ctx.withdrawer; + let lockup = lockup_for(variant, Pubkey::new_unique()); + + // Owner is not the stake program + let stake = Pubkey::new_unique(); + let stake_account = AccountSharedData::new_data_with_space( + Rent::default().minimum_balance(StakeStateV2::size_of()), + &StakeStateV2::Uninitialized, + StakeStateV2::size_of(), + &system_program_id(), + ) + .unwrap(); + + let processor: InstructionExecution = match variant { + InitializeVariant::Initialize => { + let (mut builder, accounts) = + InitializeBuilder::new().with_stake_account(&stake, &stake_account); + + let instruction = builder + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(); + + ctx.process_with(instruction, accounts) + } + InitializeVariant::InitializeChecked => { + let (mut builder, accounts) = + InitializeCheckedBuilder::new().with_stake_account(&stake, &stake_account); + + let instruction = builder + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(); + + ctx.process_with(instruction, accounts) + } + }; + + processor + .checks(&[Check::err(ProgramError::InvalidAccountOwner)]) + .test_missing_signers(false) + .execute(); +} diff --git a/program/tests/program_test.rs b/program/tests/program_test.rs index e1988c7a..4f146806 100644 --- a/program/tests/program_test.rs +++ b/program/tests/program_test.rs @@ -396,17 +396,6 @@ async fn program_test_stake_checked_instructions() { let seed = "test seed"; let seeded_address = Pubkey::create_with_seed(&seed_base, seed, &system_program::id()).unwrap(); - // Test InitializeChecked with non-signing withdrawer - let stake = create_blank_stake_account(&mut context).await; - let instruction = ixn::initialize_checked(&stake, &Authorized { staker, withdrawer }); - - process_instruction_test_missing_signers( - &mut context, - &instruction, - &vec![&withdrawer_keypair], - ) - .await; - // Test AuthorizeChecked with non-signing staker let stake = create_independent_stake_account(&mut context, &Authorized { staker, withdrawer }, 0).await; @@ -482,116 +471,6 @@ async fn program_test_stake_checked_instructions() { .await; } -#[tokio::test] -async fn program_test_stake_initialize() { - let mut context = program_test().start_with_context().await; - let accounts = Accounts::default(); - accounts.initialize(&mut context).await; - - let rent_exempt_reserve = get_stake_account_rent(&mut context.banks_client).await; - - let staker_keypair = Keypair::new(); - let withdrawer_keypair = Keypair::new(); - let custodian_keypair = Keypair::new(); - - let staker = staker_keypair.pubkey(); - let withdrawer = withdrawer_keypair.pubkey(); - let custodian = custodian_keypair.pubkey(); - - let authorized = Authorized { staker, withdrawer }; - - let lockup = Lockup { - epoch: 1, - unix_timestamp: 0, - custodian, - }; - - let stake = create_blank_stake_account(&mut context).await; - let instruction = ixn::initialize(&stake, &authorized, &lockup); - - // should pass - process_instruction(&mut context, &instruction, NO_SIGNERS) - .await - .unwrap(); - - // check that we see what we expect - let account = get_account(&mut context.banks_client, &stake).await; - let stake_state: StakeStateV2 = bincode::deserialize(&account.data).unwrap(); - assert_eq!( - stake_state, - StakeStateV2::Initialized(Meta { - authorized, - rent_exempt_reserve, - lockup, - }), - ); - - // 2nd time fails, can't move it from anything other than uninit->init - refresh_blockhash(&mut context).await; - let e = process_instruction(&mut context, &instruction, NO_SIGNERS) - .await - .unwrap_err(); - assert_eq!(e, ProgramError::InvalidAccountData); - - // not enough balance for rent - let stake = Pubkey::new_unique(); - let account = SolanaAccount { - lamports: rent_exempt_reserve / 2, - data: vec![0; StakeStateV2::size_of()], - owner: id(), - executable: false, - rent_epoch: 1000, - }; - context.set_account(&stake, &account.into()); - - let instruction = ixn::initialize(&stake, &authorized, &lockup); - let e = process_instruction(&mut context, &instruction, NO_SIGNERS) - .await - .unwrap_err(); - assert_eq!(e, ProgramError::InsufficientFunds); - - // incorrect account sizes - let stake_keypair = Keypair::new(); - let stake = stake_keypair.pubkey(); - - let instruction = system_instruction::create_account( - &context.payer.pubkey(), - &stake, - rent_exempt_reserve * 2, - StakeStateV2::size_of() as u64 + 1, - &id(), - ); - process_instruction(&mut context, &instruction, &vec![&stake_keypair]) - .await - .unwrap(); - - let instruction = ixn::initialize(&stake, &authorized, &lockup); - let e = process_instruction(&mut context, &instruction, NO_SIGNERS) - .await - .unwrap_err(); - assert_eq!(e, ProgramError::InvalidAccountData); - - let stake_keypair = Keypair::new(); - let stake = stake_keypair.pubkey(); - - let instruction = system_instruction::create_account( - &context.payer.pubkey(), - &stake, - rent_exempt_reserve, - StakeStateV2::size_of() as u64 - 1, - &id(), - ); - process_instruction(&mut context, &instruction, &vec![&stake_keypair]) - .await - .unwrap(); - - let instruction = ixn::initialize(&stake, &authorized, &lockup); - let e = process_instruction(&mut context, &instruction, NO_SIGNERS) - .await - .unwrap_err(); - assert_eq!(e, ProgramError::InvalidAccountData); -} - #[tokio::test] async fn program_test_authorize() { let mut context = program_test().start_with_context().await; From 695ce9583f8ac65ea9af869aad2b02326ce2bfd2 Mon Sep 17 00:00:00 2001 From: rustopian <96253492+rustopian@users.noreply.github.com> Date: Thu, 18 Dec 2025 18:52:44 +0000 Subject: [PATCH 2/4] simple account data provision in tests --- program/tests/helpers/context.rs | 15 +- program/tests/initialize.rs | 419 +++++++++++++------------------ 2 files changed, 186 insertions(+), 248 deletions(-) diff --git a/program/tests/helpers/context.rs b/program/tests/helpers/context.rs index ccd3faf5..db03935c 100644 --- a/program/tests/helpers/context.rs +++ b/program/tests/helpers/context.rs @@ -57,13 +57,18 @@ impl StakeTestContext { StakeAccountBuilder { lifecycle } } - /// Process an instruction - pub fn process_with( + /// Process an instruction with account data provided as a slice of (pubkey, data) pairs. + /// Sysvars are auto-resolved - only provide data for accounts that need it. + pub fn process<'b>( &self, instruction: Instruction, - accounts: Vec<(Pubkey, AccountSharedData)>, - ) -> InstructionExecution { - InstructionExecution::new(instruction, accounts, self) + accounts: &[(&Pubkey, &AccountSharedData)], + ) -> InstructionExecution<'_, 'b> { + let accounts_vec = accounts + .iter() + .map(|(pk, data)| (**pk, (*data).clone())) + .collect(); + InstructionExecution::new(instruction, accounts_vec, self) } /// Process an instruction with optional missing signer testing diff --git a/program/tests/initialize.rs b/program/tests/initialize.rs index debc92b3..c2760dde 100644 --- a/program/tests/initialize.rs +++ b/program/tests/initialize.rs @@ -3,10 +3,7 @@ mod helpers; use { - helpers::{ - context::StakeTestContext, instruction_builders::InstructionExecution, - lifecycle::StakeLifecycle, - }, + helpers::{context::StakeTestContext, lifecycle::StakeLifecycle}, mollusk_svm::result::Check, solana_account::{AccountSharedData, ReadableAccount}, solana_program_error::ProgramError, @@ -38,52 +35,6 @@ fn lockup_for( } } -/// Extension trait for InitializeBuilder to add accounts helper -trait InitializeBuilderExt { - fn with_stake_account( - self, - stake: &Pubkey, - stake_account: &AccountSharedData, - ) -> (Self, Vec<(Pubkey, AccountSharedData)>) - where - Self: Sized; -} - -impl InitializeBuilderExt for InitializeBuilder { - fn with_stake_account( - mut self, - stake: &Pubkey, - stake_account: &AccountSharedData, - ) -> (Self, Vec<(Pubkey, AccountSharedData)>) { - self.stake(*stake); - let accounts = vec![(*stake, stake_account.clone())]; - (self, accounts) - } -} - -/// Extension trait for InitializeCheckedBuilder to add accounts helper -trait InitializeCheckedBuilderExt { - fn with_stake_account( - self, - stake: &Pubkey, - stake_account: &AccountSharedData, - ) -> (Self, Vec<(Pubkey, AccountSharedData)>) - where - Self: Sized; -} - -impl InitializeCheckedBuilderExt for InitializeCheckedBuilder { - fn with_stake_account( - mut self, - stake: &Pubkey, - stake_account: &AccountSharedData, - ) -> (Self, Vec<(Pubkey, AccountSharedData)>) { - self.stake(*stake); - let accounts = vec![(*stake, stake_account.clone())]; - (self, accounts) - } -} - #[test_case(InitializeVariant::Initialize; "initialize")] #[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] fn test_initialize(variant: InitializeVariant) { @@ -107,39 +58,36 @@ fn test_initialize(variant: InitializeVariant) { .build(), ]; - let processor: InstructionExecution = match variant { - InitializeVariant::Initialize => { - let (mut builder, accounts) = - InitializeBuilder::new().with_stake_account(&stake, &stake_account); - - let instruction = builder - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(); - - ctx.process_with(instruction, accounts) - } - InitializeVariant::InitializeChecked => { - let (mut builder, accounts) = - InitializeCheckedBuilder::new().with_stake_account(&stake, &stake_account); - - let instruction = builder - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(); - - ctx.process_with(instruction, accounts) - } - }; - - processor - .checks(&checks) - .test_missing_signers(true) - .execute() + match variant { + InitializeVariant::Initialize => ctx + .process( + InitializeBuilder::new() + .stake(stake) + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(), + &[(&stake, &stake_account)], + ) + .checks(&checks) + .test_missing_signers(true) + .execute(), + InitializeVariant::InitializeChecked => ctx + .process( + InitializeCheckedBuilder::new() + .stake(stake) + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(), + &[(&stake, &stake_account)], + ) + .checks(&checks) + .test_missing_signers(true) + .execute(), + } }; let resulting_account: AccountSharedData = result.resulting_accounts[0].1.clone().into(); @@ -154,39 +102,36 @@ fn test_initialize(variant: InitializeVariant) { ); // Re-initialize should fail - let processor: InstructionExecution = match variant { - InitializeVariant::Initialize => { - let (mut builder, accounts) = - InitializeBuilder::new().with_stake_account(&stake, &resulting_account); - - let instruction = builder - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(); - - ctx.process_with(instruction, accounts) - } - InitializeVariant::InitializeChecked => { - let (mut builder, accounts) = - InitializeCheckedBuilder::new().with_stake_account(&stake, &resulting_account); - - let instruction = builder - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(); - - ctx.process_with(instruction, accounts) - } + match variant { + InitializeVariant::Initialize => ctx + .process( + InitializeBuilder::new() + .stake(stake) + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(), + &[(&stake, &resulting_account)], + ) + .checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute(), + InitializeVariant::InitializeChecked => ctx + .process( + InitializeCheckedBuilder::new() + .stake(stake) + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(), + &[(&stake, &resulting_account)], + ) + .checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute(), }; - - processor - .checks(&[Check::err(ProgramError::InvalidAccountData)]) - .test_missing_signers(false) - .execute(); } #[test_case(InitializeVariant::Initialize; "initialize")] @@ -208,39 +153,36 @@ fn test_initialize_insufficient_funds(variant: InitializeVariant) { ) .unwrap(); - let processor: InstructionExecution = match variant { - InitializeVariant::Initialize => { - let (mut builder, accounts) = - InitializeBuilder::new().with_stake_account(&stake, &stake_account); - - let instruction = builder - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(); - - ctx.process_with(instruction, accounts) - } - InitializeVariant::InitializeChecked => { - let (mut builder, accounts) = - InitializeCheckedBuilder::new().with_stake_account(&stake, &stake_account); - - let instruction = builder - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(); - - ctx.process_with(instruction, accounts) - } + match variant { + InitializeVariant::Initialize => ctx + .process( + InitializeBuilder::new() + .stake(stake) + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(), + &[(&stake, &stake_account)], + ) + .checks(&[Check::err(ProgramError::InsufficientFunds)]) + .test_missing_signers(false) + .execute(), + InitializeVariant::InitializeChecked => ctx + .process( + InitializeCheckedBuilder::new() + .stake(stake) + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(), + &[(&stake, &stake_account)], + ) + .checks(&[Check::err(ProgramError::InsufficientFunds)]) + .test_missing_signers(false) + .execute(), }; - - processor - .checks(&[Check::err(ProgramError::InsufficientFunds)]) - .test_missing_signers(false) - .execute(); } #[test_case(InitializeVariant::Initialize; "initialize")] @@ -264,39 +206,36 @@ fn test_initialize_incorrect_size_larger(variant: InitializeVariant) { ) .unwrap(); - let processor: InstructionExecution = match variant { - InitializeVariant::Initialize => { - let (mut builder, accounts) = - InitializeBuilder::new().with_stake_account(&stake, &stake_account); - - let instruction = builder - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(); - - ctx.process_with(instruction, accounts) - } - InitializeVariant::InitializeChecked => { - let (mut builder, accounts) = - InitializeCheckedBuilder::new().with_stake_account(&stake, &stake_account); - - let instruction = builder - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(); - - ctx.process_with(instruction, accounts) - } + match variant { + InitializeVariant::Initialize => ctx + .process( + InitializeBuilder::new() + .stake(stake) + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(), + &[(&stake, &stake_account)], + ) + .checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute(), + InitializeVariant::InitializeChecked => ctx + .process( + InitializeCheckedBuilder::new() + .stake(stake) + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(), + &[(&stake, &stake_account)], + ) + .checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute(), }; - - processor - .checks(&[Check::err(ProgramError::InvalidAccountData)]) - .test_missing_signers(false) - .execute(); } #[test_case(InitializeVariant::Initialize; "initialize")] @@ -320,39 +259,36 @@ fn test_initialize_incorrect_size_smaller(variant: InitializeVariant) { ) .unwrap(); - let processor: InstructionExecution = match variant { - InitializeVariant::Initialize => { - let (mut builder, accounts) = - InitializeBuilder::new().with_stake_account(&stake, &stake_account); - - let instruction = builder - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(); - - ctx.process_with(instruction, accounts) - } - InitializeVariant::InitializeChecked => { - let (mut builder, accounts) = - InitializeCheckedBuilder::new().with_stake_account(&stake, &stake_account); - - let instruction = builder - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(); - - ctx.process_with(instruction, accounts) - } + match variant { + InitializeVariant::Initialize => ctx + .process( + InitializeBuilder::new() + .stake(stake) + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(), + &[(&stake, &stake_account)], + ) + .checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute(), + InitializeVariant::InitializeChecked => ctx + .process( + InitializeCheckedBuilder::new() + .stake(stake) + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(), + &[(&stake, &stake_account)], + ) + .checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute(), }; - - processor - .checks(&[Check::err(ProgramError::InvalidAccountData)]) - .test_missing_signers(false) - .execute(); } #[test_case(InitializeVariant::Initialize; "initialize")] @@ -374,37 +310,34 @@ fn test_initialize_wrong_owner(variant: InitializeVariant) { ) .unwrap(); - let processor: InstructionExecution = match variant { - InitializeVariant::Initialize => { - let (mut builder, accounts) = - InitializeBuilder::new().with_stake_account(&stake, &stake_account); - - let instruction = builder - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(); - - ctx.process_with(instruction, accounts) - } - InitializeVariant::InitializeChecked => { - let (mut builder, accounts) = - InitializeCheckedBuilder::new().with_stake_account(&stake, &stake_account); - - let instruction = builder - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(); - - ctx.process_with(instruction, accounts) - } + match variant { + InitializeVariant::Initialize => ctx + .process( + InitializeBuilder::new() + .stake(stake) + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(), + &[(&stake, &stake_account)], + ) + .checks(&[Check::err(ProgramError::InvalidAccountOwner)]) + .test_missing_signers(false) + .execute(), + InitializeVariant::InitializeChecked => ctx + .process( + InitializeCheckedBuilder::new() + .stake(stake) + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(), + &[(&stake, &stake_account)], + ) + .checks(&[Check::err(ProgramError::InvalidAccountOwner)]) + .test_missing_signers(false) + .execute(), }; - - processor - .checks(&[Check::err(ProgramError::InvalidAccountOwner)]) - .test_missing_signers(false) - .execute(); } From edf1713901c511ca71c395c17c00ff5ea01e3d04 Mon Sep 17 00:00:00 2001 From: rustopian <96253492+rustopian@users.noreply.github.com> Date: Fri, 19 Dec 2025 11:34:51 +0000 Subject: [PATCH 3/4] better builder idiom --- program/tests/helpers/context.rs | 158 +++++++++++++----- program/tests/helpers/execution.rs | 41 +++++ program/tests/helpers/instruction_builders.rs | 61 ------- program/tests/helpers/mod.rs | 2 +- program/tests/initialize.rs | 138 +++++++-------- 5 files changed, 215 insertions(+), 185 deletions(-) create mode 100644 program/tests/helpers/execution.rs delete mode 100644 program/tests/helpers/instruction_builders.rs diff --git a/program/tests/helpers/context.rs b/program/tests/helpers/context.rs index db03935c..400a8e89 100644 --- a/program/tests/helpers/context.rs +++ b/program/tests/helpers/context.rs @@ -1,30 +1,75 @@ use { super::{ - instruction_builders::InstructionExecution, + execution::ExecutionWithChecks, lifecycle::StakeLifecycle, utils::{add_sysvars, STAKE_RENT_EXEMPTION}, }, mollusk_svm::{result::Check, Mollusk}, solana_account::AccountSharedData, solana_instruction::Instruction, + solana_program_error::ProgramError, solana_pubkey::Pubkey, + solana_stake_interface::state::Lockup, solana_stake_program::id, }; /// Builder for creating stake accounts with customizable parameters -pub struct StakeAccountBuilder { +pub struct StakeAccountBuilder<'a> { + ctx: &'a mut StakeTestContext, lifecycle: StakeLifecycle, + staked_amount: u64, + stake_authority: Option, + withdraw_authority: Option, + lockup: Option, + vote_account: Option, + stake_pubkey: Option, } -impl StakeAccountBuilder { +impl StakeAccountBuilder<'_> { + /// Set the staked amount (lamports delegated to validator) + pub fn staked_amount(mut self, amount: u64) -> Self { + self.staked_amount = amount; + self + } + + /// Set a custom stake authority (defaults to ctx.staker) + pub fn stake_authority(mut self, authority: &Pubkey) -> Self { + self.stake_authority = Some(*authority); + self + } + + /// Set a custom withdraw authority (defaults to ctx.withdrawer) + pub fn withdraw_authority(mut self, authority: &Pubkey) -> Self { + self.withdraw_authority = Some(*authority); + self + } + + /// Set a custom lockup (defaults to Lockup::default()) + pub fn lockup(mut self, lockup: &Lockup) -> Self { + self.lockup = Some(*lockup); + self + } + + /// Set a custom vote account (defaults to ctx.vote_account) + pub fn vote_account(mut self, vote_account: &Pubkey) -> Self { + self.vote_account = Some(*vote_account); + self + } + + /// Set a specific stake account pubkey (defaults to Pubkey::new_unique()) + pub fn stake_pubkey(mut self, pubkey: &Pubkey) -> Self { + self.stake_pubkey = Some(*pubkey); + self + } + + /// Build the stake account and return (pubkey, account_data) pub fn build(self) -> (Pubkey, AccountSharedData) { - let stake_pubkey = Pubkey::new_unique(); + let stake_pubkey = self.stake_pubkey.unwrap_or_else(Pubkey::new_unique); let account = self.lifecycle.create_uninitialized_account(); (stake_pubkey, account) } } -/// Consolidated test context for stake account tests pub struct StakeTestContext { pub mollusk: Mollusk, pub rent_exempt_reserve: u64, @@ -33,10 +78,8 @@ pub struct StakeTestContext { } impl StakeTestContext { - /// Create a new test context with all standard setup pub fn new() -> Self { let mollusk = Mollusk::new(&id(), "solana_stake_program"); - Self { mollusk, rent_exempt_reserve: STAKE_RENT_EXEMPTION, @@ -50,60 +93,62 @@ impl StakeTestContext { /// Example: /// ``` /// let (stake, account) = ctx - /// .stake_account(StakeLifecycle::Uninitialized) + /// .stake_account(StakeLifecycle::Active) + /// .staked_amount(1_000_000) /// .build(); /// ``` - pub fn stake_account(&mut self, lifecycle: StakeLifecycle) -> StakeAccountBuilder { - StakeAccountBuilder { lifecycle } + pub fn stake_account(&mut self, lifecycle: StakeLifecycle) -> StakeAccountBuilder<'_> { + StakeAccountBuilder { + ctx: self, + lifecycle, + staked_amount: 0, + stake_authority: None, + withdraw_authority: None, + lockup: None, + vote_account: None, + stake_pubkey: None, + } + } + + /// Configure execution with specific checks, then call .execute(instruction, accounts) + /// + /// Usage: `ctx.checks(&checks).execute(instruction, accounts)` + pub fn checks<'a, 'b>(&'a mut self, checks: &'b [Check<'b>]) -> ExecutionWithChecks<'a, 'b> { + ExecutionWithChecks::new(self, checks) } - /// Process an instruction with account data provided as a slice of (pubkey, data) pairs. - /// Sysvars are auto-resolved - only provide data for accounts that need it. - pub fn process<'b>( - &self, + /// Execute an instruction with default success checks and missing signer testing + /// + /// Usage: `ctx.execute(instruction, accounts)` + pub fn execute( + &mut self, instruction: Instruction, accounts: &[(&Pubkey, &AccountSharedData)], - ) -> InstructionExecution<'_, 'b> { - let accounts_vec = accounts - .iter() - .map(|(pk, data)| (**pk, (*data).clone())) - .collect(); - InstructionExecution::new(instruction, accounts_vec, self) + ) -> mollusk_svm::result::InstructionResult { + self.execute_internal(instruction, accounts, &[Check::success()], true) } - /// Process an instruction with optional missing signer testing - pub(crate) fn process_instruction_maybe_test_signers( - &self, - instruction: &Instruction, - accounts: Vec<(Pubkey, AccountSharedData)>, + /// Internal: execute with given checks and current config + pub(crate) fn execute_internal( + &mut self, + instruction: Instruction, + accounts: &[(&Pubkey, &AccountSharedData)], checks: &[Check], test_missing_signers: bool, ) -> mollusk_svm::result::InstructionResult { + let accounts_vec: Vec<(Pubkey, AccountSharedData)> = accounts + .iter() + .map(|(pk, data)| (**pk, (*data).clone())) + .collect(); + if test_missing_signers { - use solana_program_error::ProgramError; - - // Test that removing each signer causes failure - for i in 0..instruction.accounts.len() { - if instruction.accounts[i].is_signer { - let mut modified_instruction = instruction.clone(); - modified_instruction.accounts[i].is_signer = false; - - let accounts_with_sysvars = - add_sysvars(&self.mollusk, &modified_instruction, accounts.clone()); - - self.mollusk.process_and_validate_instruction( - &modified_instruction, - &accounts_with_sysvars, - &[Check::err(ProgramError::MissingRequiredSignature)], - ); - } - } + verify_all_signers_required(&self.mollusk, &instruction, &accounts_vec); } // Process with all signers present - let accounts_with_sysvars = add_sysvars(&self.mollusk, instruction, accounts); + let accounts_with_sysvars = add_sysvars(&self.mollusk, &instruction, accounts_vec); self.mollusk - .process_and_validate_instruction(instruction, &accounts_with_sysvars, checks) + .process_and_validate_instruction(&instruction, &accounts_with_sysvars, checks) } } @@ -112,3 +157,26 @@ impl Default for StakeTestContext { Self::new() } } + +/// Verify that removing any signer from the instruction causes MissingRequiredSignature error +fn verify_all_signers_required( + mollusk: &Mollusk, + instruction: &Instruction, + accounts: &[(Pubkey, AccountSharedData)], +) { + for i in 0..instruction.accounts.len() { + if instruction.accounts[i].is_signer { + let mut modified_instruction = instruction.clone(); + modified_instruction.accounts[i].is_signer = false; + + let accounts_with_sysvars = + add_sysvars(mollusk, &modified_instruction, accounts.to_vec()); + + mollusk.process_and_validate_instruction( + &modified_instruction, + &accounts_with_sysvars, + &[Check::err(ProgramError::MissingRequiredSignature)], + ); + } + } +} diff --git a/program/tests/helpers/execution.rs b/program/tests/helpers/execution.rs new file mode 100644 index 00000000..83e0781d --- /dev/null +++ b/program/tests/helpers/execution.rs @@ -0,0 +1,41 @@ +use { + super::context::StakeTestContext, mollusk_svm::result::Check, + solana_account::AccountSharedData, solana_instruction::Instruction, solana_pubkey::Pubkey, +}; + +/// Wrapper for executing with specific checks +/// +/// Usage: `ctx.checks(&checks).test_missing_signers(false).execute(instruction, accounts)` +pub struct ExecutionWithChecks<'a, 'b> { + pub(crate) ctx: &'a mut StakeTestContext, + pub(crate) checks: &'b [Check<'b>], + pub(crate) test_missing_signers: bool, +} + +impl<'a, 'b> ExecutionWithChecks<'a, 'b> { + pub fn new(ctx: &'a mut StakeTestContext, checks: &'b [Check<'b>]) -> Self { + Self { + ctx, + checks, + test_missing_signers: true, // default: test missing signers + } + } + + pub fn test_missing_signers(mut self, test: bool) -> Self { + self.test_missing_signers = test; + self + } + + pub fn execute( + self, + instruction: Instruction, + accounts: &[(&Pubkey, &AccountSharedData)], + ) -> mollusk_svm::result::InstructionResult { + self.ctx.execute_internal( + instruction, + accounts, + self.checks, + self.test_missing_signers, + ) + } +} diff --git a/program/tests/helpers/instruction_builders.rs b/program/tests/helpers/instruction_builders.rs deleted file mode 100644 index f805749d..00000000 --- a/program/tests/helpers/instruction_builders.rs +++ /dev/null @@ -1,61 +0,0 @@ -use { - super::context::StakeTestContext, mollusk_svm::result::Check, - solana_account::AccountSharedData, solana_instruction::Instruction, solana_pubkey::Pubkey, -}; - -/// Execution builder with validation and signer testing -pub struct InstructionExecution<'a, 'b> { - instruction: Instruction, - accounts: Vec<(Pubkey, AccountSharedData)>, - ctx: &'a StakeTestContext, - checks: Option<&'b [Check<'b>]>, - test_missing_signers: Option, // `None` runs if `Check::success` -} - -impl<'b> InstructionExecution<'_, 'b> { - pub fn checks(mut self, checks: &'b [Check<'b>]) -> Self { - self.checks = Some(checks); - self - } - - pub fn test_missing_signers(mut self, test: bool) -> Self { - self.test_missing_signers = Some(test); - self - } - - /// Executes the instruction. If `checks` is `None` or empty, uses `Check::success()`. - /// Fail-safe default: when `test_missing_signers` is `None`, runs the missing-signers - /// test (`true`). Callers must explicitly opt out with `.test_missing_signers(false)`. - pub fn execute(self) -> mollusk_svm::result::InstructionResult { - let default_checks = [Check::success()]; - let checks = match self.checks { - Some(c) if !c.is_empty() => c, - _ => &default_checks, - }; - - let test_missing_signers = self.test_missing_signers.unwrap_or(true); - - self.ctx.process_instruction_maybe_test_signers( - &self.instruction, - self.accounts, - checks, - test_missing_signers, - ) - } -} - -impl<'a> InstructionExecution<'a, '_> { - pub(crate) fn new( - instruction: Instruction, - accounts: Vec<(Pubkey, AccountSharedData)>, - ctx: &'a StakeTestContext, - ) -> Self { - Self { - instruction, - accounts, - ctx, - checks: None, - test_missing_signers: None, - } - } -} diff --git a/program/tests/helpers/mod.rs b/program/tests/helpers/mod.rs index ff09b000..adbdb63a 100644 --- a/program/tests/helpers/mod.rs +++ b/program/tests/helpers/mod.rs @@ -2,6 +2,6 @@ #![allow(dead_code)] pub mod context; -pub mod instruction_builders; +pub mod execution; pub mod lifecycle; pub mod utils; diff --git a/program/tests/initialize.rs b/program/tests/initialize.rs index c2760dde..57edb929 100644 --- a/program/tests/initialize.rs +++ b/program/tests/initialize.rs @@ -38,7 +38,7 @@ fn lockup_for( #[test_case(InitializeVariant::Initialize; "initialize")] #[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] fn test_initialize(variant: InitializeVariant) { - let mut ctx = StakeTestContext::new(); + let mut ctx = StakeTestContext::default(); let staker = ctx.staker; let withdrawer = ctx.withdrawer; @@ -59,34 +59,26 @@ fn test_initialize(variant: InitializeVariant) { ]; match variant { - InitializeVariant::Initialize => ctx - .process( - InitializeBuilder::new() - .stake(stake) - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(), - &[(&stake, &stake_account)], - ) - .checks(&checks) - .test_missing_signers(true) - .execute(), - InitializeVariant::InitializeChecked => ctx - .process( - InitializeCheckedBuilder::new() - .stake(stake) - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(), - &[(&stake, &stake_account)], - ) - .checks(&checks) - .test_missing_signers(true) - .execute(), + InitializeVariant::Initialize => ctx.checks(&checks).execute( + InitializeBuilder::new() + .stake(stake) + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(), + &[(&stake, &stake_account)], + ), + InitializeVariant::InitializeChecked => ctx.checks(&checks).execute( + InitializeCheckedBuilder::new() + .stake(stake) + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(), + &[(&stake, &stake_account)], + ), } }; @@ -104,7 +96,9 @@ fn test_initialize(variant: InitializeVariant) { // Re-initialize should fail match variant { InitializeVariant::Initialize => ctx - .process( + .checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute( InitializeBuilder::new() .stake(stake) .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) @@ -115,29 +109,25 @@ fn test_initialize(variant: InitializeVariant) { }) .instruction(), &[(&stake, &resulting_account)], - ) + ), + InitializeVariant::InitializeChecked => ctx .checks(&[Check::err(ProgramError::InvalidAccountData)]) .test_missing_signers(false) - .execute(), - InitializeVariant::InitializeChecked => ctx - .process( + .execute( InitializeCheckedBuilder::new() .stake(stake) .stake_authority(staker) .withdraw_authority(withdrawer) .instruction(), &[(&stake, &resulting_account)], - ) - .checks(&[Check::err(ProgramError::InvalidAccountData)]) - .test_missing_signers(false) - .execute(), + ), }; } #[test_case(InitializeVariant::Initialize; "initialize")] #[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] fn test_initialize_insufficient_funds(variant: InitializeVariant) { - let ctx = StakeTestContext::new(); + let mut ctx = StakeTestContext::default(); let staker = ctx.staker; let withdrawer = ctx.withdrawer; @@ -155,7 +145,9 @@ fn test_initialize_insufficient_funds(variant: InitializeVariant) { match variant { InitializeVariant::Initialize => ctx - .process( + .checks(&[Check::err(ProgramError::InsufficientFunds)]) + .test_missing_signers(false) + .execute( InitializeBuilder::new() .stake(stake) .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) @@ -166,29 +158,25 @@ fn test_initialize_insufficient_funds(variant: InitializeVariant) { }) .instruction(), &[(&stake, &stake_account)], - ) + ), + InitializeVariant::InitializeChecked => ctx .checks(&[Check::err(ProgramError::InsufficientFunds)]) .test_missing_signers(false) - .execute(), - InitializeVariant::InitializeChecked => ctx - .process( + .execute( InitializeCheckedBuilder::new() .stake(stake) .stake_authority(staker) .withdraw_authority(withdrawer) .instruction(), &[(&stake, &stake_account)], - ) - .checks(&[Check::err(ProgramError::InsufficientFunds)]) - .test_missing_signers(false) - .execute(), + ), }; } #[test_case(InitializeVariant::Initialize; "initialize")] #[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] fn test_initialize_incorrect_size_larger(variant: InitializeVariant) { - let ctx = StakeTestContext::new(); + let mut ctx = StakeTestContext::default(); let rent_exempt_reserve = Rent::default().minimum_balance(StakeStateV2::size_of() * 2); @@ -208,7 +196,9 @@ fn test_initialize_incorrect_size_larger(variant: InitializeVariant) { match variant { InitializeVariant::Initialize => ctx - .process( + .checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute( InitializeBuilder::new() .stake(stake) .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) @@ -219,29 +209,25 @@ fn test_initialize_incorrect_size_larger(variant: InitializeVariant) { }) .instruction(), &[(&stake, &stake_account)], - ) + ), + InitializeVariant::InitializeChecked => ctx .checks(&[Check::err(ProgramError::InvalidAccountData)]) .test_missing_signers(false) - .execute(), - InitializeVariant::InitializeChecked => ctx - .process( + .execute( InitializeCheckedBuilder::new() .stake(stake) .stake_authority(staker) .withdraw_authority(withdrawer) .instruction(), &[(&stake, &stake_account)], - ) - .checks(&[Check::err(ProgramError::InvalidAccountData)]) - .test_missing_signers(false) - .execute(), + ), }; } #[test_case(InitializeVariant::Initialize; "initialize")] #[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] fn test_initialize_incorrect_size_smaller(variant: InitializeVariant) { - let ctx = StakeTestContext::new(); + let mut ctx = StakeTestContext::default(); let rent_exempt_reserve = Rent::default().minimum_balance(StakeStateV2::size_of()); @@ -261,7 +247,9 @@ fn test_initialize_incorrect_size_smaller(variant: InitializeVariant) { match variant { InitializeVariant::Initialize => ctx - .process( + .checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute( InitializeBuilder::new() .stake(stake) .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) @@ -272,29 +260,25 @@ fn test_initialize_incorrect_size_smaller(variant: InitializeVariant) { }) .instruction(), &[(&stake, &stake_account)], - ) + ), + InitializeVariant::InitializeChecked => ctx .checks(&[Check::err(ProgramError::InvalidAccountData)]) .test_missing_signers(false) - .execute(), - InitializeVariant::InitializeChecked => ctx - .process( + .execute( InitializeCheckedBuilder::new() .stake(stake) .stake_authority(staker) .withdraw_authority(withdrawer) .instruction(), &[(&stake, &stake_account)], - ) - .checks(&[Check::err(ProgramError::InvalidAccountData)]) - .test_missing_signers(false) - .execute(), + ), }; } #[test_case(InitializeVariant::Initialize; "initialize")] #[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] fn test_initialize_wrong_owner(variant: InitializeVariant) { - let ctx = StakeTestContext::new(); + let mut ctx = StakeTestContext::default(); let staker = ctx.staker; let withdrawer = ctx.withdrawer; @@ -312,7 +296,9 @@ fn test_initialize_wrong_owner(variant: InitializeVariant) { match variant { InitializeVariant::Initialize => ctx - .process( + .checks(&[Check::err(ProgramError::InvalidAccountOwner)]) + .test_missing_signers(false) + .execute( InitializeBuilder::new() .stake(stake) .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) @@ -323,21 +309,17 @@ fn test_initialize_wrong_owner(variant: InitializeVariant) { }) .instruction(), &[(&stake, &stake_account)], - ) + ), + InitializeVariant::InitializeChecked => ctx .checks(&[Check::err(ProgramError::InvalidAccountOwner)]) .test_missing_signers(false) - .execute(), - InitializeVariant::InitializeChecked => ctx - .process( + .execute( InitializeCheckedBuilder::new() .stake(stake) .stake_authority(staker) .withdraw_authority(withdrawer) .instruction(), &[(&stake, &stake_account)], - ) - .checks(&[Check::err(ProgramError::InvalidAccountOwner)]) - .test_missing_signers(false) - .execute(), + ), }; } From 080449a6107697b85c060ce403416ea6426e63ad Mon Sep 17 00:00:00 2001 From: rustopian <96253492+rustopian@users.noreply.github.com> Date: Fri, 19 Dec 2025 12:02:22 +0000 Subject: [PATCH 4/4] nice little helper --- program/tests/initialize.rs | 233 ++++++++++-------------------------- 1 file changed, 64 insertions(+), 169 deletions(-) diff --git a/program/tests/initialize.rs b/program/tests/initialize.rs index 57edb929..9920d801 100644 --- a/program/tests/initialize.rs +++ b/program/tests/initialize.rs @@ -6,12 +6,13 @@ use { helpers::{context::StakeTestContext, lifecycle::StakeLifecycle}, mollusk_svm::result::Check, solana_account::{AccountSharedData, ReadableAccount}, + solana_instruction::Instruction, solana_program_error::ProgramError, solana_pubkey::Pubkey, solana_rent::Rent, solana_sdk_ids::{stake::id, system_program::id as system_program_id}, solana_stake_client::instructions::{InitializeBuilder, InitializeCheckedBuilder}, - solana_stake_interface::state::StakeStateV2, + solana_stake_interface::state::{Lockup, StakeStateV2}, test_case::test_case, }; @@ -21,17 +22,39 @@ enum InitializeVariant { InitializeChecked, } -fn lockup_for( - variant: InitializeVariant, - custodian: Pubkey, -) -> solana_stake_interface::state::Lockup { +fn lockup_for(variant: InitializeVariant, custodian: Pubkey) -> Lockup { match variant { - InitializeVariant::Initialize => solana_stake_interface::state::Lockup { + InitializeVariant::Initialize => Lockup { epoch: 1, unix_timestamp: 0, custodian, }, - InitializeVariant::InitializeChecked => solana_stake_interface::state::Lockup::default(), + InitializeVariant::InitializeChecked => Lockup::default(), + } +} + +fn build_initialize_instruction( + variant: InitializeVariant, + stake: Pubkey, + staker: Pubkey, + withdrawer: Pubkey, + lockup: Lockup, +) -> Instruction { + match variant { + InitializeVariant::Initialize => InitializeBuilder::new() + .stake(stake) + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(), + InitializeVariant::InitializeChecked => InitializeCheckedBuilder::new() + .stake(stake) + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(), } } @@ -58,28 +81,10 @@ fn test_initialize(variant: InitializeVariant) { .build(), ]; - match variant { - InitializeVariant::Initialize => ctx.checks(&checks).execute( - InitializeBuilder::new() - .stake(stake) - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(), - &[(&stake, &stake_account)], - ), - InitializeVariant::InitializeChecked => ctx.checks(&checks).execute( - InitializeCheckedBuilder::new() - .stake(stake) - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(), - &[(&stake, &stake_account)], - ), - } + ctx.checks(&checks).execute( + build_initialize_instruction(variant, stake, staker, withdrawer, lockup), + &[(&stake, &stake_account)], + ) }; let resulting_account: AccountSharedData = result.resulting_accounts[0].1.clone().into(); @@ -94,34 +99,12 @@ fn test_initialize(variant: InitializeVariant) { ); // Re-initialize should fail - match variant { - InitializeVariant::Initialize => ctx - .checks(&[Check::err(ProgramError::InvalidAccountData)]) - .test_missing_signers(false) - .execute( - InitializeBuilder::new() - .stake(stake) - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(), - &[(&stake, &resulting_account)], - ), - InitializeVariant::InitializeChecked => ctx - .checks(&[Check::err(ProgramError::InvalidAccountData)]) - .test_missing_signers(false) - .execute( - InitializeCheckedBuilder::new() - .stake(stake) - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(), - &[(&stake, &resulting_account)], - ), - }; + ctx.checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute( + build_initialize_instruction(variant, stake, staker, withdrawer, lockup), + &[(&stake, &resulting_account)], + ); } #[test_case(InitializeVariant::Initialize; "initialize")] @@ -143,34 +126,12 @@ fn test_initialize_insufficient_funds(variant: InitializeVariant) { ) .unwrap(); - match variant { - InitializeVariant::Initialize => ctx - .checks(&[Check::err(ProgramError::InsufficientFunds)]) - .test_missing_signers(false) - .execute( - InitializeBuilder::new() - .stake(stake) - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(), - &[(&stake, &stake_account)], - ), - InitializeVariant::InitializeChecked => ctx - .checks(&[Check::err(ProgramError::InsufficientFunds)]) - .test_missing_signers(false) - .execute( - InitializeCheckedBuilder::new() - .stake(stake) - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(), - &[(&stake, &stake_account)], - ), - }; + ctx.checks(&[Check::err(ProgramError::InsufficientFunds)]) + .test_missing_signers(false) + .execute( + build_initialize_instruction(variant, stake, staker, withdrawer, lockup), + &[(&stake, &stake_account)], + ); } #[test_case(InitializeVariant::Initialize; "initialize")] @@ -194,34 +155,12 @@ fn test_initialize_incorrect_size_larger(variant: InitializeVariant) { ) .unwrap(); - match variant { - InitializeVariant::Initialize => ctx - .checks(&[Check::err(ProgramError::InvalidAccountData)]) - .test_missing_signers(false) - .execute( - InitializeBuilder::new() - .stake(stake) - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(), - &[(&stake, &stake_account)], - ), - InitializeVariant::InitializeChecked => ctx - .checks(&[Check::err(ProgramError::InvalidAccountData)]) - .test_missing_signers(false) - .execute( - InitializeCheckedBuilder::new() - .stake(stake) - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(), - &[(&stake, &stake_account)], - ), - }; + ctx.checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute( + build_initialize_instruction(variant, stake, staker, withdrawer, lockup), + &[(&stake, &stake_account)], + ); } #[test_case(InitializeVariant::Initialize; "initialize")] @@ -245,34 +184,12 @@ fn test_initialize_incorrect_size_smaller(variant: InitializeVariant) { ) .unwrap(); - match variant { - InitializeVariant::Initialize => ctx - .checks(&[Check::err(ProgramError::InvalidAccountData)]) - .test_missing_signers(false) - .execute( - InitializeBuilder::new() - .stake(stake) - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(), - &[(&stake, &stake_account)], - ), - InitializeVariant::InitializeChecked => ctx - .checks(&[Check::err(ProgramError::InvalidAccountData)]) - .test_missing_signers(false) - .execute( - InitializeCheckedBuilder::new() - .stake(stake) - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(), - &[(&stake, &stake_account)], - ), - }; + ctx.checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute( + build_initialize_instruction(variant, stake, staker, withdrawer, lockup), + &[(&stake, &stake_account)], + ); } #[test_case(InitializeVariant::Initialize; "initialize")] @@ -294,32 +211,10 @@ fn test_initialize_wrong_owner(variant: InitializeVariant) { ) .unwrap(); - match variant { - InitializeVariant::Initialize => ctx - .checks(&[Check::err(ProgramError::InvalidAccountOwner)]) - .test_missing_signers(false) - .execute( - InitializeBuilder::new() - .stake(stake) - .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) - .arg1(solana_stake_client::types::Lockup { - unix_timestamp: lockup.unix_timestamp, - epoch: lockup.epoch, - custodian: lockup.custodian, - }) - .instruction(), - &[(&stake, &stake_account)], - ), - InitializeVariant::InitializeChecked => ctx - .checks(&[Check::err(ProgramError::InvalidAccountOwner)]) - .test_missing_signers(false) - .execute( - InitializeCheckedBuilder::new() - .stake(stake) - .stake_authority(staker) - .withdraw_authority(withdrawer) - .instruction(), - &[(&stake, &stake_account)], - ), - }; + ctx.checks(&[Check::err(ProgramError::InvalidAccountOwner)]) + .test_missing_signers(false) + .execute( + build_initialize_instruction(variant, stake, staker, withdrawer, lockup), + &[(&stake, &stake_account)], + ); }