bonfida_test_utils/
program_test_context_ext.rs

1use async_trait::async_trait;
2use solana_program::{
3    clock::Clock, instruction::Instruction, program_pack::Pack, pubkey::Pubkey, rent::Rent,
4    system_instruction::create_account,
5};
6use solana_program_test::ProgramTestContext;
7use solana_sdk::{signature::Keypair, signer::Signer, transaction::Transaction};
8
9use crate::error::TestError;
10
11const NANOSECONDS_IN_SECOND: u128 = 1_000_000_000;
12
13#[async_trait]
14pub trait ProgramTestContextExt {
15    async fn mint_tokens(
16        &mut self,
17        mint_authority: &Keypair,
18        mint_pubkey: &Pubkey,
19        token_account: &Pubkey,
20        amount: u64,
21    ) -> Result<(), TestError>;
22
23    async fn get_token_account(
24        &mut self,
25        key: Pubkey,
26    ) -> Result<spl_token::state::Account, TestError>;
27
28    async fn sign_send_instructions(
29        &mut self,
30        instructions: &[Instruction],
31        signers: &[&Keypair],
32    ) -> Result<(), TestError>;
33
34    async fn warp_to_timestamp(&mut self, timestamp: i64) -> Result<(), TestError>;
35    async fn initialize_token_accounts(
36        &mut self,
37        mint: Pubkey,
38        owners: &[Pubkey],
39    ) -> Result<Vec<Pubkey>, TestError>;
40
41    async fn get_current_timestamp(&mut self) -> Result<i64, TestError>;
42
43    async fn initialize_new_account(
44        &mut self,
45        space: usize,
46        program_id: Pubkey,
47    ) -> Result<Pubkey, TestError>;
48}
49
50#[async_trait]
51impl ProgramTestContextExt for ProgramTestContext {
52    async fn mint_tokens(
53        &mut self,
54        mint_authority: &Keypair,
55        mint_pubkey: &Pubkey,
56        token_account: &Pubkey,
57        amount: u64,
58    ) -> Result<(), TestError> {
59        let mint_instruction = spl_token::instruction::mint_to(
60            &spl_token::ID,
61            mint_pubkey,
62            token_account,
63            &mint_authority.pubkey(),
64            &[],
65            amount,
66        )?;
67        self.sign_send_instructions(&[mint_instruction], &[mint_authority])
68            .await?;
69        Ok(())
70    }
71    async fn get_token_account(
72        &mut self,
73        key: Pubkey,
74    ) -> Result<spl_token::state::Account, TestError> {
75        let raw_account = self
76            .banks_client
77            .get_account(key)
78            .await?
79            .ok_or(TestError::AccountDoesNotExist)?;
80        if raw_account.owner != spl_token::ID {
81            return Err(TestError::InvalidTokenAccount);
82        }
83        Ok(spl_token::state::Account::unpack(&raw_account.data)?)
84    }
85    async fn sign_send_instructions(
86        &mut self,
87        instructions: &[Instruction],
88        signers: &[&Keypair],
89    ) -> Result<(), TestError> {
90        let mut transaction = Transaction::new_with_payer(instructions, Some(&self.payer.pubkey()));
91        let mut payer_signers = Vec::with_capacity(1 + signers.len());
92        payer_signers.push(&self.payer);
93        for s in signers {
94            payer_signers.push(s);
95        }
96        transaction.partial_sign(&payer_signers, self.last_blockhash);
97        self.banks_client.process_transaction(transaction).await?;
98        Ok(())
99    }
100
101    async fn warp_to_timestamp(&mut self, timestamp: i64) -> Result<(), TestError> {
102        let mut clock = self.banks_client.get_sysvar::<Clock>().await?;
103        if clock.unix_timestamp > timestamp {
104            return Err(TestError::InvalidTimestampForWarp);
105        }
106        let ns_per_slot = self.genesis_config().ns_per_slot();
107        let time_delta_ns =
108            (timestamp - clock.unix_timestamp).unsigned_abs() as u128 * NANOSECONDS_IN_SECOND;
109        let number_of_slots_to_warp: u64 = (time_delta_ns / ns_per_slot).try_into().unwrap();
110
111        clock.unix_timestamp = timestamp;
112        self.set_sysvar(&clock);
113        self.warp_to_slot(clock.slot + number_of_slots_to_warp)?;
114        Ok(())
115    }
116
117    async fn initialize_token_accounts(
118        &mut self,
119        mint: Pubkey,
120        owners: &[Pubkey],
121    ) -> Result<Vec<Pubkey>, TestError> {
122        let mut instructions = Vec::with_capacity(owners.len());
123        let mut account_keys = Vec::with_capacity(owners.len());
124        for o in owners {
125            let i = spl_associated_token_account::instruction::create_associated_token_account(
126                &self.payer.pubkey(),
127                o,
128                &mint,
129                &spl_token::ID,
130            );
131            account_keys.push(i.accounts[1].pubkey);
132            instructions.push(i);
133        }
134        for c in instructions.chunks(10) {
135            self.sign_send_instructions(c, &[]).await?;
136        }
137        Ok(account_keys)
138    }
139
140    async fn get_current_timestamp(&mut self) -> Result<i64, TestError> {
141        let clock = self.banks_client.get_sysvar::<Clock>().await?;
142        Ok(clock.unix_timestamp)
143    }
144
145    async fn initialize_new_account(
146        &mut self,
147        space: usize,
148        program_id: Pubkey,
149    ) -> Result<Pubkey, TestError> {
150        let account_keypair = Keypair::new();
151        let rent = self.banks_client.get_sysvar::<Rent>().await.unwrap();
152        let lamports = rent.minimum_balance(space);
153        let ix = create_account(
154            &self.payer.pubkey(),
155            &account_keypair.pubkey(),
156            lamports,
157            space as u64,
158            &program_id,
159        );
160        self.sign_send_instructions(&[ix], &[&account_keypair])
161            .await?;
162        Ok(account_keypair.pubkey())
163    }
164}
bonfida_test_utils/program_test_context_ext.rs

bonfida_test_utils/
program_test_context_ext.rs
