migrated Initialize and InitializeChecked tests to mollusk

StakeTestContext to reduce boilerplate

InstructionConfig builder pattern

improve InitializeChecked coverage

Author: rustopian

Cargo.lock

@@ -45,6 +45,7 @@ solana-system-interface = { version = "2.0.0", features = ["bincode"] }
 solana-sysvar-id = "3.0.0"
 solana-transaction = "3.0.0"
 test-case = "3.3.1"
+tokio = { version = "1", features = ["full"] }
 
 [lib]
 crate-type = ["cdylib", "lib"]

program/Cargo.toml

@@ -0,0 +1,112 @@
+use {
+    super::{
+        instruction_builders::{InstructionConfig, 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<'b, C: InstructionConfig>(
+        &self,
+        config: C,
+    ) -> InstructionExecution<'_, 'b> {
+        InstructionExecution::new(
+            config.build_instruction(self),
+            config.build_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()
+    }
+}

program/tests/helpers/context.rs

@@ -0,0 +1,96 @@
+use {
+    super::context::StakeTestContext,
+    mollusk_svm::result::Check,
+    solana_account::AccountSharedData,
+    solana_instruction::Instruction,
+    solana_pubkey::Pubkey,
+    solana_stake_interface::{
+        instruction as ixn,
+        state::{Authorized, Lockup},
+    },
+};
+
+// Trait for instruction configuration that builds instruction and accounts
+pub trait InstructionConfig {
+    fn build_instruction(&self, ctx: &StakeTestContext) -> Instruction;
+    fn build_accounts(&self) -> Vec<(Pubkey, AccountSharedData)>;
+}
+
+/// 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: bool,
+}
+
+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) -> Self {
+        self.test_missing_signers = true;
+        self
+    }
+
+    pub fn execute(self) -> mollusk_svm::result::InstructionResult {
+        let default_checks = [Check::success()];
+        let checks = self.checks.unwrap_or(&default_checks);
+        self.ctx.process_instruction_maybe_test_signers(
+            &self.instruction,
+            self.accounts,
+            checks,
+            self.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: false,
+        }
+    }
+}
+
+pub struct InitializeConfig<'a> {
+    pub stake: (&'a Pubkey, &'a AccountSharedData),
+    pub authorized: &'a Authorized,
+    pub lockup: &'a Lockup,
+}
+
+impl InstructionConfig for InitializeConfig<'_> {
+    fn build_instruction(&self, _ctx: &StakeTestContext) -> Instruction {
+        ixn::initialize(self.stake.0, self.authorized, self.lockup)
+    }
+
+    fn build_accounts(&self) -> Vec<(Pubkey, AccountSharedData)> {
+        vec![(*self.stake.0, self.stake.1.clone())]
+    }
+}
+
+pub struct InitializeCheckedConfig<'a> {
+    pub stake: (&'a Pubkey, &'a AccountSharedData),
+    pub authorized: &'a Authorized,
+}
+
+impl InstructionConfig for InitializeCheckedConfig<'_> {
+    fn build_instruction(&self, _ctx: &StakeTestContext) -> Instruction {
+        ixn::initialize_checked(self.stake.0, self.authorized)
+    }
+
+    fn build_accounts(&self) -> Vec<(Pubkey, AccountSharedData)> {
+        vec![(*self.stake.0, self.stake.1.clone())]
+    }
+}

program/tests/helpers/instruction_builders.rs

@@ -0,0 +1,30 @@
+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
+#[allow(dead_code)]
+#[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()
+    }
+}

program/tests/helpers/lifecycle.rs

@@ -0,0 +1,10 @@
+#![allow(clippy::arithmetic_side_effects)]
+
+pub mod context;
+pub mod instruction_builders;
+pub mod lifecycle;
+pub mod utils;
+
+pub use {
+    context::StakeTestContext, instruction_builders::InitializeConfig, lifecycle::StakeLifecycle,
+};

program/tests/helpers/mod.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<Pubkey, Account> = 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
+}