better builder idiom

Author: rustopian

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<Pubkey>,
+    withdraw_authority: Option<Pubkey>,
+    lockup: Option<Lockup>,
+    vote_account: Option<Pubkey>,
+    stake_pubkey: Option<Pubkey>,
 }
 
-impl StakeAccountBuilder {
+impl<'a> StakeAccountBuilder<'a> {
+    /// 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)],
+            );
+        }
+    }
+}

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,
+        )
+    }
+}

program/tests/helpers/instruction_builders.rs

@@ -2,6 +2,6 @@
 #![allow(dead_code)]
 
 pub mod context;
-pub mod instruction_builders;
+pub mod execution;
 pub mod lifecycle;
 pub mod utils;