mollusk_helper/
context.rs

1use crate::account;
2use crate::account_store::InMemoryAccountStore;
3use crate::error::{MolluskHelperError, Result};
4use crate::token;
5use crate::transaction::TransactionBuilder;
6use mollusk_svm::account_store::AccountStore;
7use mollusk_svm::result::{InstructionResult, ProgramResult};
8use mollusk_svm::{Mollusk, MolluskContext};
9use solana_account::Account;
10use solana_address::Address;
11use solana_hash::Hash;
12use solana_instruction::Instruction;
13use solana_keypair::Keypair;
14use solana_message::{v0, AddressLookupTableAccount, VersionedMessage};
15use solana_program_pack::Pack;
16use solana_pubkey::Pubkey;
17use solana_signer::Signer;
18use spl_token::state::Account as TokenAccount;
19use std::collections::HashMap;
20use std::sync::{Arc, RwLock};
21
22pub const MEMO_PROGRAM_ID: Pubkey =
23    solana_pubkey::pubkey!("MemoSq4gqABAXKb96qnH8TysNcWxMyWCqXgDLGmfcHr");
24
25pub const MEMO_V1_PROGRAM_ID: Pubkey =
26    solana_pubkey::pubkey!("Memo1UhkJRfHyvLMcVucJwxXeuD728EqVDDwQDxFMNo");
27
28pub const TOKEN_2022_PROGRAM_ID: Pubkey =
29    solana_pubkey::pubkey!("TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb");
30
31pub const ADDRESS_LOOKUP_TABLE_PROGRAM_ID: Pubkey =
32    solana_pubkey::pubkey!("AddressLookupTab1e1111111111111111111111111");
33
34pub const COMPUTE_BUDGET_PROGRAM_ID: Pubkey =
35    solana_pubkey::pubkey!("ComputeBudget111111111111111111111111111111");
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum ProgramLoader {
39    V2,
40    V3,
41}
42
43fn add_default_programs(mollusk: &mut Mollusk) {
44    mollusk_svm_programs_token::token::add_program(mollusk);
45    mollusk_svm_programs_token::associated_token::add_program(mollusk);
46    mollusk_svm_programs_token::token2022::add_program(mollusk);
47    mollusk_svm_programs_memo::memo::add_program(mollusk);
48    mollusk_svm_programs_memo::memo_v1::add_program(mollusk);
49}
50
51pub struct MolluskContextHelper {
52    pub(crate) context: MolluskContext<InMemoryAccountStore>,
53    keypairs: Arc<RwLock<HashMap<String, Keypair>>>,
54}
55
56impl MolluskContextHelper {
57    pub fn new(program_id: &Pubkey, elf_bytes: &[u8]) -> Self {
58        Self::new_with_options(
59            program_id,
60            elf_bytes,
61            ProgramLoader::V3,
62            Self::current_unix_timestamp(),
63        )
64    }
65
66    pub fn new_with_timestamp(program_id: &Pubkey, elf_bytes: &[u8], unix_timestamp: u64) -> Self {
67        Self::new_with_options(program_id, elf_bytes, ProgramLoader::V3, unix_timestamp)
68    }
69
70    pub fn new_with_loader(program_id: &Pubkey, elf_bytes: &[u8], loader: ProgramLoader) -> Self {
71        Self::new_with_options(
72            program_id,
73            elf_bytes,
74            loader,
75            Self::current_unix_timestamp(),
76        )
77    }
78
79    pub fn new_with_options(
80        program_id: &Pubkey,
81        elf_bytes: &[u8],
82        loader: ProgramLoader,
83        unix_timestamp: u64,
84    ) -> Self {
85        let mut mollusk = Mollusk::default();
86
87        let loader_key = match loader {
88            ProgramLoader::V2 => &mollusk_svm::program::loader_keys::LOADER_V2,
89            ProgramLoader::V3 => &mollusk_svm::program::loader_keys::LOADER_V3,
90        };
91
92        mollusk.add_program_with_loader_and_elf(
93            &Self::pubkey_to_address(program_id),
94            loader_key,
95            elf_bytes,
96        );
97
98        add_default_programs(&mut mollusk);
99        mollusk.sysvars.clock.unix_timestamp = unix_timestamp as i64;
100
101        let store = InMemoryAccountStore::new();
102        let context = mollusk.with_context(store);
103
104        Self {
105            context,
106            keypairs: Arc::new(RwLock::new(HashMap::new())),
107        }
108    }
109
110    pub fn new_without_program() -> Self {
111        Self::new_without_program_with_timestamp(Self::current_unix_timestamp())
112    }
113
114    pub fn new_without_program_with_timestamp(unix_timestamp: u64) -> Self {
115        let mut mollusk = Mollusk::default();
116
117        add_default_programs(&mut mollusk);
118        mollusk.sysvars.clock.unix_timestamp = unix_timestamp as i64;
119
120        let store = InMemoryAccountStore::new();
121        let context = mollusk.with_context(store);
122
123        Self {
124            context,
125            keypairs: Arc::new(RwLock::new(HashMap::new())),
126        }
127    }
128
129    pub fn add_program(&mut self, program_id: &Pubkey, elf_bytes: &[u8]) {
130        self.add_program_with_loader(program_id, elf_bytes, ProgramLoader::V3);
131    }
132
133    pub fn add_program_with_loader(
134        &mut self,
135        program_id: &Pubkey,
136        elf_bytes: &[u8],
137        loader: ProgramLoader,
138    ) {
139        let loader_key = match loader {
140            ProgramLoader::V2 => &mollusk_svm::program::loader_keys::LOADER_V2,
141            ProgramLoader::V3 => &mollusk_svm::program::loader_keys::LOADER_V3,
142        };
143
144        self.context.mollusk.add_program_with_loader_and_elf(
145            &Self::pubkey_to_address(program_id),
146            loader_key,
147            elf_bytes,
148        );
149    }
150
151    pub fn current_unix_timestamp() -> u64 {
152        std::time::SystemTime::now()
153            .duration_since(std::time::UNIX_EPOCH)
154            .expect("Time went backwards")
155            .as_secs()
156    }
157
158    pub fn process_instruction(&self, instruction: &Instruction) -> Result<InstructionResult> {
159        let result = self.context.process_instruction(instruction);
160        match &result.program_result {
161            ProgramResult::Success => Ok(result),
162            ProgramResult::Failure(e) => Err(MolluskHelperError::ProgramError(e.clone())),
163            ProgramResult::UnknownError(e) => Err(MolluskHelperError::InstructionFailed(e.clone())),
164        }
165    }
166
167    pub fn process_instruction_unchecked(&self, instruction: &Instruction) -> InstructionResult {
168        self.context.process_instruction(instruction)
169    }
170
171    pub(crate) fn process_instruction_internal(
172        &self,
173        instruction: &Instruction,
174    ) -> InstructionResult {
175        self.context.process_instruction(instruction)
176    }
177
178    pub fn transaction(&self) -> TransactionBuilder<'_> {
179        TransactionBuilder::new(self)
180    }
181
182    pub(crate) fn snapshot_accounts(&self) -> HashMap<Address, Account> {
183        self.context.account_store.borrow().snapshot()
184    }
185
186    pub(crate) fn restore_accounts(&self, snapshot: HashMap<Address, Account>) {
187        self.context.account_store.borrow_mut().restore(snapshot);
188    }
189
190    pub fn add_account(&self, pubkey: &Pubkey, account: Account) {
191        let address = Self::pubkey_to_address(pubkey);
192        self.context
193            .account_store
194            .borrow_mut()
195            .add_account(address, account);
196    }
197
198    pub fn get_account(&self, pubkey: &Pubkey) -> Option<Account> {
199        let address = Self::pubkey_to_address(pubkey);
200        self.context.account_store.borrow().get_account(&address)
201    }
202
203    pub fn get_balance(&self, pubkey: &Pubkey) -> Option<u64> {
204        let address = Self::pubkey_to_address(pubkey);
205        self.context.account_store.borrow().get_balance(&address)
206    }
207
208    pub fn fund_account(&self, pubkey: &Pubkey, lamports: u64) {
209        let account = account::system_account_with_lamports(lamports);
210        self.add_account(pubkey, account);
211    }
212
213    pub fn store_keypair(&self, name: &str, keypair: Keypair) -> Result<()> {
214        self.keypairs
215            .write()
216            .map_err(|_| MolluskHelperError::LockError)?
217            .insert(name.to_string(), keypair);
218        Ok(())
219    }
220
221    pub fn sign_with(&self, name: &str, message: &[u8]) -> Result<[u8; 64]> {
222        let keypairs = self
223            .keypairs
224            .read()
225            .map_err(|_| MolluskHelperError::LockError)?;
226        let keypair = keypairs
227            .get(name)
228            .ok_or_else(|| MolluskHelperError::KeypairNotFound(name.to_string()))?;
229        Ok(keypair.sign_message(message).into())
230    }
231
232    pub fn get_keypair_pubkey(&self, name: &str) -> Result<Pubkey> {
233        let keypairs = self
234            .keypairs
235            .read()
236            .map_err(|_| MolluskHelperError::LockError)?;
237        let keypair = keypairs
238            .get(name)
239            .ok_or_else(|| MolluskHelperError::KeypairNotFound(name.to_string()))?;
240        Ok(keypair.pubkey())
241    }
242
243    pub fn update_unix_timestamp(&mut self, timestamp: i64) {
244        self.context.mollusk.sysvars.clock.unix_timestamp = timestamp;
245    }
246
247    pub fn get_unix_timestamp(&self) -> i64 {
248        self.context.mollusk.sysvars.clock.unix_timestamp
249    }
250
251    pub fn warp_to_slot(&mut self, slot: u64) {
252        self.context.mollusk.warp_to_slot(slot);
253    }
254
255    pub fn create_mint(&self, mint_pubkey: &Pubkey, authority: &Pubkey, decimals: u8) {
256        let account = token::create_mint_account(authority, decimals);
257        self.add_account(mint_pubkey, account);
258    }
259
260    pub fn create_token_account(
261        &self,
262        token_account_pubkey: &Pubkey,
263        mint: &Pubkey,
264        owner: &Pubkey,
265        amount: u64,
266    ) {
267        let account = token::create_token_account(mint, owner, amount);
268        self.add_account(token_account_pubkey, account);
269    }
270
271    pub fn create_native_token_account(
272        &self,
273        token_account_pubkey: &Pubkey,
274        owner: &Pubkey,
275        lamports: u64,
276    ) {
277        let account = token::create_native_token_account(owner, lamports);
278        self.add_account(token_account_pubkey, account);
279    }
280
281    pub fn get_token_balance(&self, token_account_pubkey: &Pubkey) -> Result<u64> {
282        let account = self
283            .get_account(token_account_pubkey)
284            .ok_or_else(|| MolluskHelperError::AccountNotFound(token_account_pubkey.to_string()))?;
285        let token_account =
286            TokenAccount::unpack(&account.data).map_err(MolluskHelperError::ProgramError)?;
287        Ok(token_account.amount)
288    }
289
290    pub fn mint_to(
291        &self,
292        mint: &Pubkey,
293        destination: &Pubkey,
294        authority: &Pubkey,
295        amount: u64,
296    ) -> Result<InstructionResult> {
297        let ix = token::mint_to_instruction(mint, destination, authority, amount);
298        self.process_instruction(&ix)
299    }
300
301    pub fn transfer_tokens(
302        &self,
303        source: &Pubkey,
304        destination: &Pubkey,
305        authority: &Pubkey,
306        amount: u64,
307    ) -> Result<InstructionResult> {
308        let ix = token::transfer_instruction(source, destination, authority, amount);
309        self.process_instruction(&ix)
310    }
311
312    pub fn sync_native(&self, token_account: &Pubkey) -> Result<InstructionResult> {
313        let ix = token::sync_native_instruction(token_account);
314        self.process_instruction(&ix)
315    }
316
317    pub fn transfer_sol(
318        &self,
319        from: &Pubkey,
320        to: &Pubkey,
321        lamports: u64,
322    ) -> Result<InstructionResult> {
323        let ix = solana_system_interface::instruction::transfer(from, to, lamports);
324        self.process_instruction(&ix)
325    }
326
327    pub fn get_associated_token_address(&self, wallet: &Pubkey, mint: &Pubkey) -> Pubkey {
328        token::get_associated_token_address(wallet, mint)
329    }
330
331    pub fn create_lookup_table_account(
332        key: Pubkey,
333        addresses: Vec<Pubkey>,
334    ) -> AddressLookupTableAccount {
335        AddressLookupTableAccount { key, addresses }
336    }
337
338    pub fn build_versioned_tx_size(
339        payer: &Pubkey,
340        instructions: &[Instruction],
341        lookup_tables: &[AddressLookupTableAccount],
342    ) -> (usize, usize, i64) {
343        let recent_blockhash = Hash::default();
344        let message =
345            v0::Message::try_compile(payer, instructions, lookup_tables, recent_blockhash)
346                .expect("Failed to compile v0 message");
347        let versioned_message = VersionedMessage::V0(message);
348
349        let num_signers = instructions
350            .iter()
351            .flat_map(|ix| ix.accounts.iter())
352            .filter(|a| a.is_signer)
353            .map(|a| a.pubkey)
354            .collect::<std::collections::HashSet<_>>()
355            .len()
356            .max(1);
357
358        let message_bytes =
359            bincode::serialize(&versioned_message).expect("Failed to serialize message");
360        let tx_size = 1 + (num_signers * 64) + message_bytes.len();
361        let tx_limit: usize = 1232;
362        let remaining = tx_limit as i64 - tx_size as i64;
363
364        (tx_size, tx_limit, remaining)
365    }
366
367    pub fn pubkey_to_address(pubkey: &Pubkey) -> Address {
368        Address::new_from_array(pubkey.to_bytes())
369    }
370
371    pub fn address_to_pubkey(address: &Address) -> Pubkey {
372        Pubkey::from(address.to_bytes())
373    }
374
375    pub fn system_program() -> Pubkey {
376        solana_pubkey::pubkey!("11111111111111111111111111111111")
377    }
378
379    pub fn token_program() -> Pubkey {
380        token::TOKEN_PROGRAM_ID
381    }
382
383    pub fn token_2022_program() -> Pubkey {
384        TOKEN_2022_PROGRAM_ID
385    }
386
387    pub fn associated_token_program() -> Pubkey {
388        token::ASSOCIATED_TOKEN_PROGRAM_ID
389    }
390
391    pub fn memo_program() -> Pubkey {
392        MEMO_PROGRAM_ID
393    }
394
395    pub fn memo_v1_program() -> Pubkey {
396        MEMO_V1_PROGRAM_ID
397    }
398
399    pub fn address_lookup_table_program() -> Pubkey {
400        ADDRESS_LOOKUP_TABLE_PROGRAM_ID
401    }
402
403    pub fn compute_budget_program() -> Pubkey {
404        COMPUTE_BUDGET_PROGRAM_ID
405    }
406
407    pub fn native_mint() -> Pubkey {
408        token::NATIVE_MINT
409    }
410
411    pub fn rent_sysvar() -> Pubkey {
412        solana_pubkey::pubkey!("SysvarRent111111111111111111111111111111111")
413    }
414
415    pub fn add_program_account(&self, pubkey: &Pubkey, owner: &Pubkey, data: Vec<u8>) {
416        let account = account::program_account(owner, data);
417        self.add_account(pubkey, account);
418    }
419
420    pub fn create_associated_token_account_instruction(
421        payer: &Pubkey,
422        wallet: &Pubkey,
423        mint: &Pubkey,
424    ) -> Instruction {
425        token::create_associated_token_account_instruction(payer, wallet, mint)
426    }
427}
mollusk_helper/context.rs

mollusk_helper/
context.rs
