Program tests

.gitignore

@@ -1,6 +1,7 @@
 node_modules
 dist
 ts/**/*.js
+.vscode
 
 .anchor
 target

programs/mango-v3-reimbursement/Cargo.toml

@@ -14,9 +14,22 @@ no-idl = []
 no-log-ix-name = []
 cpi = ["no-entrypoint"]
 default = []
+test-bpf = []
 
 [dependencies]
 anchor-lang = { path = "../../../mango-v4/anchor/lang", features = ["init-if-needed"] }
 anchor-spl = { path = "../../../mango-v4/anchor/spl" }
 solana-program = "~1.10.29"
-static_assertions = "1.1"
+static_assertions = "1.1"
+
+[dev-dependencies]
+solana-sdk = { version = "~1.10.29", default-features = false }
+solana-program-test = "~1.10.29"
+solana-logger = "~1.10.29"
+async-trait = "0.1.52"
+bytemuck = "^1.7.2"
+spl-token = { version = "^3.0.0", features = ["no-entrypoint"] }
+spl-associated-token-account = { version = "^1.0.3", features = ["no-entrypoint"] }
+log = "0.4.14"
+env_logger = "0.9.0"
+lazy_static = "1.4.0"

programs/mango-v3-reimbursement/tests/program_test/client.rs

@@ -0,0 +1,361 @@
+#![allow(dead_code)]
+
+use anchor_lang::prelude::*;
+use anchor_lang::solana_program::sysvar;
+use anchor_spl::token::Token;
+use solana_program_test::BanksClientError;
+use solana_sdk::instruction;
+use solana_sdk::transport::TransportError;
+use std::sync::Arc;
+
+use super::solana::SolanaCookie;
+use super::utils::TestKeypair;
+use mango_v3_reimbursement::state::*;
+
+#[async_trait::async_trait(?Send)]
+pub trait ClientAccountLoader {
+    async fn load_bytes(&self, pubkey: &Pubkey) -> Option<Vec<u8>>;
+    async fn load<T: AccountDeserialize>(&self, pubkey: &Pubkey) -> Option<T> {
+        let bytes = self.load_bytes(pubkey).await?;
+        AccountDeserialize::try_deserialize(&mut &bytes[..]).ok()
+    }
+}
+
+#[async_trait::async_trait(?Send)]
+impl ClientAccountLoader for &SolanaCookie {
+    async fn load_bytes(&self, pubkey: &Pubkey) -> Option<Vec<u8>> {
+        self.get_account_data(*pubkey).await
+    }
+}
+
+// TODO: report error outwards etc
+pub async fn send_tx<CI: ClientInstruction>(
+    solana: &SolanaCookie,
+    ix: CI,
+) -> std::result::Result<CI::Accounts, TransportError> {
+    let (accounts, instruction) = ix.to_instruction(solana).await;
+    let signers = ix.signers();
+    let instructions = vec![instruction];
+    solana
+        .process_transaction(&instructions, Some(&signers[..]))
+        .await?;
+    Ok(accounts)
+}
+
+/// Build a transaction from multiple instructions
+pub struct ClientTransaction {
+    solana: Arc<SolanaCookie>,
+    instructions: Vec<instruction::Instruction>,
+    signers: Vec<TestKeypair>,
+}
+
+impl<'a> ClientTransaction {
+    pub fn new(solana: &Arc<SolanaCookie>) -> Self {
+        Self {
+            solana: solana.clone(),
+            instructions: vec![],
+            signers: vec![],
+        }
+    }
+
+    pub async fn add_instruction<CI: ClientInstruction>(&mut self, ix: CI) -> CI::Accounts {
+        let solana: &SolanaCookie = &self.solana;
+        let (accounts, instruction) = ix.to_instruction(solana).await;
+        self.instructions.push(instruction);
+        self.signers.extend(ix.signers());
+        accounts
+    }
+
+    pub fn add_instruction_direct(&mut self, ix: instruction::Instruction) {
+        self.instructions.push(ix);
+    }
+
+    pub fn add_signer(&mut self, keypair: TestKeypair) {
+        self.signers.push(keypair);
+    }
+
+    pub async fn send(&self) -> std::result::Result<(), BanksClientError> {
+        self.solana
+            .process_transaction(&self.instructions, Some(&self.signers))
+            .await
+    }
+}
+
+#[async_trait::async_trait(?Send)]
+pub trait ClientInstruction {
+    type Accounts: anchor_lang::ToAccountMetas;
+    type Instruction: anchor_lang::InstructionData;
+
+    async fn to_instruction(
+        &self,
+        loader: impl ClientAccountLoader + 'async_trait,
+    ) -> (Self::Accounts, instruction::Instruction);
+    fn signers(&self) -> Vec<TestKeypair>;
+}
+
+fn make_instruction(
+    program_id: Pubkey,
+    accounts: &impl anchor_lang::ToAccountMetas,
+    data: impl anchor_lang::InstructionData,
+) -> instruction::Instruction {
+    instruction::Instruction {
+        program_id,
+        accounts: anchor_lang::ToAccountMetas::to_account_metas(accounts, None),
+        data: anchor_lang::InstructionData::data(&data),
+    }
+}
+
+//
+// a struct for each instruction along with its
+// ClientInstruction impl
+//
+
+pub struct CreateGroupInstruction {
+    pub group_num: u32,
+    pub claim_transfer_destination: Pubkey,
+    pub testing: bool,
+    pub table: Pubkey,
+    pub payer: TestKeypair,
+    pub authority: TestKeypair,
+}
+#[async_trait::async_trait(?Send)]
+impl ClientInstruction for CreateGroupInstruction {
+    type Accounts = mango_v3_reimbursement::accounts::CreateGroup;
+    type Instruction = mango_v3_reimbursement::instruction::CreateGroup;
+    async fn to_instruction(
+        &self,
+        _account_loader: impl ClientAccountLoader + 'async_trait,
+    ) -> (Self::Accounts, instruction::Instruction) {
+        let program_id = mango_v3_reimbursement::id();
+        let instruction = Self::Instruction {
+            group_num: self.group_num,
+            claim_transfer_destination: self.claim_transfer_destination,
+            testing: if self.testing { 1 } else { 0 },
+        };
+
+        let group = Pubkey::find_program_address(
+            &[b"Group".as_ref(), &self.group_num.to_le_bytes()],
+            &program_id,
+        )
+        .0;
+
+        let accounts = Self::Accounts {
+            group,
+            table: self.table,
+            payer: self.payer.pubkey(),
+            authority: self.authority.pubkey(),
+            system_program: System::id(),
+        };
+
+        let instruction = make_instruction(program_id, &accounts, instruction);
+        (accounts, instruction)
+    }
+
+    fn signers(&self) -> Vec<TestKeypair> {
+        vec![self.payer, self.authority]
+    }
+}
+
+pub struct CreateVaultInstruction {
+    pub group: Pubkey,
+    pub authority: TestKeypair,
+    pub token_index: usize,
+    pub mint: Pubkey,
+    pub payer: TestKeypair,
+}
+#[async_trait::async_trait(?Send)]
+impl ClientInstruction for CreateVaultInstruction {
+    type Accounts = mango_v3_reimbursement::accounts::CreateVault;
+    type Instruction = mango_v3_reimbursement::instruction::CreateVault;
+    async fn to_instruction(
+        &self,
+        account_loader: impl ClientAccountLoader + 'async_trait,
+    ) -> (Self::Accounts, instruction::Instruction) {
+        let program_id = mango_v3_reimbursement::id();
+
+        let group: Group = account_loader.load(&self.group).await.unwrap();
+
+        let instruction = Self::Instruction {
+            token_index: self.token_index,
+        };
+
+        let claim_mint = Pubkey::find_program_address(
+            &[
+                b"Mint".as_ref(),
+                self.group.as_ref(),
+                &self.token_index.to_le_bytes(),
+            ],
+            &program_id,
+        )
+        .0;
+
+        let claim_transfer_token_account =
+            anchor_spl::associated_token::get_associated_token_address(
+                &group.claim_transfer_destination,
+                &claim_mint,
+            );
+
+        let vault =
+            anchor_spl::associated_token::get_associated_token_address(&self.group, &self.mint);
+
+        let accounts = Self::Accounts {
+            group: self.group,
+            vault,
+            mint: self.mint,
+            authority: self.authority.pubkey(),
+            payer: self.payer.pubkey(),
+            claim_transfer_token_account,
+            claim_transfer_destination: group.claim_transfer_destination,
+            claim_mint,
+            token_program: Token::id(),
+            system_program: System::id(),
+            rent: sysvar::rent::id(),
+            associated_token_program: anchor_spl::associated_token::ID,
+        };
+
+        let instruction = make_instruction(program_id, &accounts, instruction);
+        (accounts, instruction)
+    }
+
+    fn signers(&self) -> Vec<TestKeypair> {
+        vec![self.authority, self.payer]
+    }
+}
+
+pub struct StartReimbursementInstruction {
+    pub group: Pubkey,
+    pub authority: TestKeypair,
+}
+#[async_trait::async_trait(?Send)]
+impl ClientInstruction for StartReimbursementInstruction {
+    type Accounts = mango_v3_reimbursement::accounts::StartReimbursement;
+    type Instruction = mango_v3_reimbursement::instruction::StartReimbursement;
+    async fn to_instruction(
+        &self,
+        _account_loader: impl ClientAccountLoader + 'async_trait,
+    ) -> (Self::Accounts, instruction::Instruction) {
+        let program_id = mango_v3_reimbursement::id();
+        let instruction = Self::Instruction {};
+
+        let accounts = Self::Accounts {
+            group: self.group,
+            authority: self.authority.pubkey(),
+        };
+
+        let instruction = make_instruction(program_id, &accounts, instruction);
+        (accounts, instruction)
+    }
+
+    fn signers(&self) -> Vec<TestKeypair> {
+        vec![self.authority]
+    }
+}
+
+pub struct CreateReimbursementAccountInstruction {
+    pub group: Pubkey,
+    pub mango_account_owner: Pubkey,
+    pub payer: TestKeypair,
+}
+#[async_trait::async_trait(?Send)]
+impl ClientInstruction for CreateReimbursementAccountInstruction {
+    type Accounts = mango_v3_reimbursement::accounts::CreateReimbursementAccount;
+    type Instruction = mango_v3_reimbursement::instruction::CreateReimbursementAccount;
+    async fn to_instruction(
+        &self,
+        _account_loader: impl ClientAccountLoader + 'async_trait,
+    ) -> (Self::Accounts, instruction::Instruction) {
+        let program_id = mango_v3_reimbursement::id();
+        let instruction = Self::Instruction {};
+
+        let reimbursement_account = Pubkey::find_program_address(
+            &[
+                b"ReimbursementAccount".as_ref(),
+                self.group.as_ref(),
+                self.mango_account_owner.as_ref(),
+            ],
+            &program_id,
+        )
+        .0;
+
+        let accounts = Self::Accounts {
+            group: self.group,
+            reimbursement_account,
+            mango_account_owner: self.mango_account_owner,
+            payer: self.payer.pubkey(),
+            system_program: System::id(),
+            rent: sysvar::rent::id(),
+        };
+
+        let instruction = make_instruction(program_id, &accounts, instruction);
+        (accounts, instruction)
+    }
+
+    fn signers(&self) -> Vec<TestKeypair> {
+        vec![self.payer]
+    }
+}
+
+pub struct ReimburseInstruction {
+    pub group: Pubkey,
+    pub token_index: usize,
+    pub mango_account_owner: TestKeypair,
+    pub transfer_claim: bool,
+    pub token_account: Pubkey,
+    pub table_index: usize,
+}
+#[async_trait::async_trait(?Send)]
+impl ClientInstruction for ReimburseInstruction {
+    type Accounts = mango_v3_reimbursement::accounts::Reimburse;
+    type Instruction = mango_v3_reimbursement::instruction::Reimburse;
+    async fn to_instruction(
+        &self,
+        account_loader: impl ClientAccountLoader + 'async_trait,
+    ) -> (Self::Accounts, instruction::Instruction) {
+        let program_id = mango_v3_reimbursement::id();
+
+        let group: Group = account_loader.load(&self.group).await.unwrap();
+
+        let instruction = Self::Instruction {
+            index_into_table: self.table_index,
+            token_index: self.token_index,
+            transfer_claim: self.transfer_claim,
+        };
+
+        let reimbursement_account = Pubkey::find_program_address(
+            &[
+                b"ReimbursementAccount".as_ref(),
+                self.group.as_ref(),
+                self.mango_account_owner.pubkey().as_ref(),
+            ],
+            &program_id,
+        )
+        .0;
+
+        let claim_mint_token_account = anchor_spl::associated_token::get_associated_token_address(
+            &group.claim_transfer_destination,
+            &group.claim_mints[self.token_index],
+        );
+
+        let accounts = Self::Accounts {
+            group: self.group,
+            vault: group.vaults[self.token_index],
+            token_account: self.token_account,
+            reimbursement_account,
+            mango_account_owner: self.mango_account_owner.pubkey(),
+            signer: self.mango_account_owner.pubkey(),
+            claim_mint_token_account,
+            claim_mint: group.claim_mints[self.token_index],
+            table: group.table,
+            token_program: Token::id(),
+            system_program: System::id(),
+            rent: sysvar::rent::id(),
+        };
+
+        let instruction = make_instruction(program_id, &accounts, instruction);
+        (accounts, instruction)
+    }
+
+    fn signers(&self) -> Vec<TestKeypair> {
+        vec![self.mango_account_owner]
+    }
+}

programs/mango-v3-reimbursement/tests/program_test/cookies.rs

@@ -0,0 +1,20 @@
+use solana_program::pubkey::*;
+
+use crate::utils::*;
+
+#[derive(Debug, Clone, Copy)]
+pub struct MintCookie {
+    pub index: usize,
+    pub decimals: u8,
+    pub unit: f64,
+    pub base_lot: f64,
+    pub quote_lot: f64,
+    pub pubkey: Pubkey,
+    pub authority: TestKeypair,
+}
+
+#[derive(Debug, Clone)]
+pub struct UserCookie {
+    pub key: TestKeypair,
+    pub token_accounts: Vec<Pubkey>,
+}

programs/mango-v3-reimbursement/tests/program_test/mod.rs

@@ -0,0 +1,252 @@
+#![allow(dead_code)]
+
+use std::cell::RefCell;
+use std::{sync::Arc, sync::RwLock};
+
+use log::*;
+use solana_program::{program_option::COption, program_pack::Pack};
+use solana_program_test::*;
+use solana_sdk::pubkey::Pubkey;
+use spl_token::{state::*, *};
+
+pub use client::*;
+pub use cookies::*;
+pub use solana::*;
+pub use utils::*;
+
+pub mod client;
+pub mod cookies;
+pub mod solana;
+pub mod utils;
+
+pub trait AddPacked {
+    fn add_packable_account<T: Pack>(
+        &mut self,
+        pubkey: Pubkey,
+        amount: u64,
+        data: &T,
+        owner: &Pubkey,
+    );
+}
+
+impl AddPacked for ProgramTest {
+    fn add_packable_account<T: Pack>(
+        &mut self,
+        pubkey: Pubkey,
+        amount: u64,
+        data: &T,
+        owner: &Pubkey,
+    ) {
+        let mut account = solana_sdk::account::Account::new(amount, T::get_packed_len(), owner);
+        data.pack_into_slice(&mut account.data);
+        self.add_account(pubkey, account);
+    }
+}
+
+struct LoggerWrapper {
+    inner: env_logger::Logger,
+    capture: Arc<RwLock<Vec<String>>>,
+}
+
+impl Log for LoggerWrapper {
+    fn enabled(&self, metadata: &log::Metadata) -> bool {
+        self.inner.enabled(metadata)
+    }
+
+    fn log(&self, record: &log::Record) {
+        if record
+            .target()
+            .starts_with("solana_runtime::message_processor")
+        {
+            let msg = record.args().to_string();
+            if let Some(data) = msg.strip_prefix("Program log: ") {
+                self.capture.write().unwrap().push(data.into());
+            } else if let Some(data) = msg.strip_prefix("Program data: ") {
+                self.capture.write().unwrap().push(data.into());
+            }
+        }
+        self.inner.log(record);
+    }
+
+    fn flush(&self) {}
+}
+
+pub struct TestContextBuilder {
+    test: ProgramTest,
+    logger_capture: Arc<RwLock<Vec<String>>>,
+    mint0: Pubkey,
+}
+
+lazy_static::lazy_static! {
+    static ref LOGGER_CAPTURE: Arc<RwLock<Vec<String>>> = Arc::new(RwLock::new(vec![]));
+    static ref LOGGER_LOCK: Arc<RwLock<()>> = Arc::new(RwLock::new(()));
+}
+
+impl TestContextBuilder {
+    pub fn new() -> Self {
+        // We need to intercept logs to capture program log output
+        let log_filter = "solana_rbpf=trace,\
+                    solana_runtime::message_processor=debug,\
+                    solana_runtime::system_instruction_processor=trace,\
+                    solana_program_test=info";
+        let env_logger =
+            env_logger::Builder::from_env(env_logger::Env::new().default_filter_or(log_filter))
+                .format_timestamp_nanos()
+                .build();
+        let _ = log::set_boxed_logger(Box::new(LoggerWrapper {
+            inner: env_logger,
+            capture: LOGGER_CAPTURE.clone(),
+        }));
+
+        let mut test = ProgramTest::new(
+            "mango_v3_reimbursement",
+            mango_v3_reimbursement::id(),
+            processor!(mango_v3_reimbursement::entry),
+        );
+
+        // intentionally set to as tight as possible, to catch potential problems early
+        test.set_compute_max_units(200000);
+
+        Self {
+            test,
+            logger_capture: LOGGER_CAPTURE.clone(),
+            mint0: Pubkey::new_unique(),
+        }
+    }
+
+    pub fn test(&mut self) -> &mut ProgramTest {
+        &mut self.test
+    }
+
+    pub fn create_mints(&mut self) -> Vec<MintCookie> {
+        let mut mints: Vec<MintCookie> = vec![
+            MintCookie {
+                index: 0,
+                decimals: 6,
+                unit: 10u64.pow(6) as f64,
+                base_lot: 100 as f64,
+                quote_lot: 10 as f64,
+                pubkey: self.mint0,
+                authority: TestKeypair::new(),
+            }, // symbol: "MNGO".to_string()
+        ];
+        for i in 1..16 {
+            mints.push(MintCookie {
+                index: i,
+                decimals: 6,
+                unit: 10u64.pow(6) as f64,
+                base_lot: 0 as f64,
+                quote_lot: 0 as f64,
+                pubkey: Pubkey::default(),
+                authority: TestKeypair::new(),
+            });
+        }
+        // Add mints in loop
+        for mint_index in 0..mints.len() {
+            let mint_pk: Pubkey;
+            if mints[mint_index].pubkey == Pubkey::default() {
+                mint_pk = Pubkey::new_unique();
+            } else {
+                mint_pk = mints[mint_index].pubkey;
+            }
+            mints[mint_index].pubkey = mint_pk;
+
+            self.test.add_packable_account(
+                mint_pk,
+                u32::MAX as u64,
+                &Mint {
+                    is_initialized: true,
+                    mint_authority: COption::Some(mints[mint_index].authority.pubkey()),
+                    decimals: mints[mint_index].decimals,
+                    ..Mint::default()
+                },
+                &spl_token::id(),
+            );
+        }
+
+        mints
+    }
+
+    pub fn create_users(&mut self, mints: &[MintCookie]) -> Vec<UserCookie> {
+        let num_users = 4;
+        let mut users = Vec::new();
+        for _ in 0..num_users {
+            let user_key = TestKeypair::new();
+            self.test.add_account(
+                user_key.pubkey(),
+                solana_sdk::account::Account::new(
+                    u32::MAX as u64,
+                    0,
+                    &solana_sdk::system_program::id(),
+                ),
+            );
+
+            // give every user 10^18 (< 2^60) of every token
+            // ~~ 1 trillion in case of 6 decimals
+            let mut token_accounts = Vec::new();
+            for mint_index in 0..mints.len() {
+                let token_key = Pubkey::new_unique();
+                self.test.add_packable_account(
+                    token_key,
+                    u32::MAX as u64,
+                    &spl_token::state::Account {
+                        mint: mints[mint_index].pubkey,
+                        owner: user_key.pubkey(),
+                        amount: 1_000_000_000_000_000_000,
+                        state: spl_token::state::AccountState::Initialized,
+                        ..spl_token::state::Account::default()
+                    },
+                    &spl_token::id(),
+                );
+
+                token_accounts.push(token_key);
+            }
+            users.push(UserCookie {
+                key: user_key,
+                token_accounts,
+            });
+        }
+
+        users
+    }
+
+    pub async fn start_default(mut self) -> TestContext {
+        let mints = self.create_mints();
+        let users = self.create_users(&mints);
+
+        let solana = self.start().await;
+
+        TestContext {
+            solana: solana.clone(),
+            mints,
+            users,
+        }
+    }
+
+    pub async fn start(self) -> Arc<SolanaCookie> {
+        let mut context = self.test.start_with_context().await;
+        let rent = context.banks_client.get_rent().await.unwrap();
+
+        let solana = Arc::new(SolanaCookie {
+            context: RefCell::new(context),
+            rent,
+            logger_capture: self.logger_capture.clone(),
+            logger_lock: LOGGER_LOCK.clone(),
+            last_transaction_log: RefCell::new(vec![]),
+        });
+
+        solana
+    }
+}
+
+pub struct TestContext {
+    pub solana: Arc<SolanaCookie>,
+    pub mints: Vec<MintCookie>,
+    pub users: Vec<UserCookie>,
+}
+
+impl TestContext {
+    pub async fn new() -> Self {
+        TestContextBuilder::new().start_default().await
+    }
+}

programs/mango-v3-reimbursement/tests/program_test/solana.rs

@@ -0,0 +1,225 @@
+#![allow(dead_code)]
+
+use std::cell::RefCell;
+use std::sync::{Arc, RwLock};
+
+use crate::utils::TestKeypair;
+use anchor_lang::AccountDeserialize;
+use anchor_spl::token::TokenAccount;
+use solana_program::{program_pack::Pack, rent::*, system_instruction};
+use solana_program_test::*;
+use solana_sdk::{
+    account::ReadableAccount,
+    instruction::Instruction,
+    pubkey::Pubkey,
+    signature::{Keypair, Signer},
+    transaction::Transaction,
+};
+//use spl_token::*;
+
+pub struct SolanaCookie {
+    pub context: RefCell<ProgramTestContext>,
+    pub rent: Rent,
+    pub logger_capture: Arc<RwLock<Vec<String>>>,
+    pub logger_lock: Arc<RwLock<()>>,
+    pub last_transaction_log: RefCell<Vec<String>>,
+}
+
+impl SolanaCookie {
+    pub async fn process_transaction(
+        &self,
+        instructions: &[Instruction],
+        signers: Option<&[TestKeypair]>,
+    ) -> Result<(), BanksClientError> {
+        // The locking in this function is convoluted:
+        // We capture the program log output by overriding the global logger and capturing
+        // messages there. This logger is potentially shared among multiple tests that run
+        // concurrently.
+        // To allow each independent SolanaCookie to capture only the logs from the transaction
+        // passed to process_transaction, wo globally hold the "program_log_lock" for the
+        // duration that the tx needs to process. So only a single one can run at a time.
+        let tx_log_lock = Arc::new(self.logger_lock.write().unwrap());
+        self.logger_capture.write().unwrap().clear();
+
+        let mut context = self.context.borrow_mut();
+
+        let mut transaction =
+            Transaction::new_with_payer(&instructions, Some(&context.payer.pubkey()));
+
+        let mut all_signers = vec![&context.payer];
+        let signer_keypairs =
+            signers.map(|signers| signers.iter().map(|s| s.into()).collect::<Vec<Keypair>>());
+        let signer_keypair_refs = signer_keypairs
+            .as_ref()
+            .map(|kps| kps.iter().map(|kp| kp).collect::<Vec<&Keypair>>());
+
+        if let Some(signer_keypair_refs) = signer_keypair_refs {
+            all_signers.extend(signer_keypair_refs.iter());
+        }
+
+        // This fails when warping is involved - https://gitmemory.com/issue/solana-labs/solana/18201/868325078
+        // let recent_blockhash = self.context.banks_client.get_recent_blockhash().await.unwrap();
+
+        transaction.sign(&all_signers, context.last_blockhash);
+
+        let result = context
+            .banks_client
+            .process_transaction_with_commitment(
+                transaction,
+                solana_sdk::commitment_config::CommitmentLevel::Processed,
+            )
+            .await;
+
+        *self.last_transaction_log.borrow_mut() = self.logger_capture.read().unwrap().clone();
+
+        drop(tx_log_lock);
+        drop(context);
+
+        // This makes sure every transaction gets a new blockhash, avoiding issues where sending
+        // the same transaction again would lead to it being skipped.
+        self.advance_by_slots(1).await;
+
+        result
+    }
+
+    pub async fn get_clock(&self) -> solana_program::clock::Clock {
+        self.context
+            .borrow_mut()
+            .banks_client
+            .get_sysvar::<solana_program::clock::Clock>()
+            .await
+            .unwrap()
+    }
+
+    pub async fn advance_by_slots(&self, slots: u64) {
+        let clock = self.get_clock().await;
+        self.context
+            .borrow_mut()
+            .warp_to_slot(clock.slot + slots + 1)
+            .unwrap();
+    }
+
+    pub async fn advance_clock(&self) {
+        let mut clock = self.get_clock().await;
+        let old_ts = clock.unix_timestamp;
+
+        // just advance enough to ensure we get changes over last_updated in various ix
+        // if this gets too slow for our tests, remove and replace with manual time offset
+        // which is configurable
+        while clock.unix_timestamp <= old_ts {
+            self.context
+                .borrow_mut()
+                .warp_to_slot(clock.slot + 50)
+                .unwrap();
+            clock = self.get_clock().await;
+        }
+    }
+
+    pub async fn get_newest_slot_from_history(&self) -> u64 {
+        self.context
+            .borrow_mut()
+            .banks_client
+            .get_sysvar::<solana_program::slot_history::SlotHistory>()
+            .await
+            .unwrap()
+            .newest()
+    }
+
+    pub async fn create_account_from_len(&self, owner: &Pubkey, len: usize) -> Pubkey {
+        let key = TestKeypair::new();
+        let rent = self.rent.minimum_balance(len);
+        let create_account_instr = solana_sdk::system_instruction::create_account(
+            &self.context.borrow().payer.pubkey(),
+            &key.pubkey(),
+            rent,
+            len as u64,
+            &owner,
+        );
+        self.process_transaction(&[create_account_instr], Some(&[key]))
+            .await
+            .unwrap();
+        key.pubkey()
+    }
+
+    pub async fn create_account_for_type<T>(&self, owner: &Pubkey) -> Pubkey {
+        let key = TestKeypair::new();
+        let len = 8 + std::mem::size_of::<T>();
+        let rent = self.rent.minimum_balance(len);
+        let create_account_instr = solana_sdk::system_instruction::create_account(
+            &self.context.borrow().payer.pubkey(),
+            &key.pubkey(),
+            rent,
+            len as u64,
+            &owner,
+        );
+        self.process_transaction(&[create_account_instr], Some(&[key]))
+            .await
+            .unwrap();
+        key.pubkey()
+    }
+
+    pub async fn create_token_account(&self, owner: &Pubkey, mint: Pubkey) -> Pubkey {
+        let keypair = TestKeypair::new();
+        let rent = self.rent.minimum_balance(spl_token::state::Account::LEN);
+
+        let instructions = [
+            system_instruction::create_account(
+                &self.context.borrow().payer.pubkey(),
+                &keypair.pubkey(),
+                rent,
+                spl_token::state::Account::LEN as u64,
+                &spl_token::id(),
+            ),
+            spl_token::instruction::initialize_account(
+                &spl_token::id(),
+                &keypair.pubkey(),
+                &mint,
+                owner,
+            )
+            .unwrap(),
+        ];
+
+        self.process_transaction(&instructions, Some(&[keypair]))
+            .await
+            .unwrap();
+        return keypair.pubkey();
+    }
+
+    pub async fn get_account_data(&self, address: Pubkey) -> Option<Vec<u8>> {
+        Some(
+            self.context
+                .borrow_mut()
+                .banks_client
+                .get_account(address)
+                .await
+                .unwrap()?
+                .data()
+                .to_vec(),
+        )
+    }
+
+    pub async fn get_account_opt<T: AccountDeserialize>(&self, address: Pubkey) -> Option<T> {
+        let data = self.get_account_data(address).await?;
+        let mut data_slice: &[u8] = &data;
+        AccountDeserialize::try_deserialize(&mut data_slice).ok()
+    }
+
+    // Use when accounts are too big for the stack
+    pub async fn get_account_boxed<T: AccountDeserialize>(&self, address: Pubkey) -> Box<T> {
+        let data = self.get_account_data(address).await.unwrap();
+        let mut data_slice: &[u8] = &data;
+        Box::new(AccountDeserialize::try_deserialize(&mut data_slice).unwrap())
+    }
+
+    pub async fn get_account<T: AccountDeserialize>(&self, address: Pubkey) -> T {
+        self.get_account_opt(address).await.unwrap()
+    }
+
+    pub async fn token_account_balance(&self, address: Pubkey) -> u64 {
+        self.get_account::<TokenAccount>(address).await.amount
+    }
+
+    pub fn program_log(&self) -> Vec<String> {
+        self.last_transaction_log.borrow().clone()
+    }
+}

programs/mango-v3-reimbursement/tests/program_test/utils.rs

@@ -0,0 +1,53 @@
+#![allow(dead_code)]
+
+use solana_sdk::pubkey::Pubkey;
+use solana_sdk::signature::Keypair;
+
+pub fn clone_keypair(keypair: &Keypair) -> Keypair {
+    Keypair::from_base58_string(&keypair.to_base58_string())
+}
+
+// Add clone() to Keypair, totally safe in tests
+pub trait ClonableKeypair {
+    fn clone(&self) -> Self;
+}
+impl ClonableKeypair for Keypair {
+    fn clone(&self) -> Self {
+        clone_keypair(self)
+    }
+}
+
+/// A Keypair-like struct that's Clone and Copy and can be into()ed to a Keypair
+///
+/// The regular Keypair is neither Clone nor Copy because the key data is sensitive
+/// and should not be copied needlessly. That just makes things difficult for tests.
+#[derive(Clone, Copy, Debug)]
+pub struct TestKeypair([u8; 64]);
+impl TestKeypair {
+    pub fn new() -> Self {
+        Keypair::new().into()
+    }
+
+    pub fn to_keypair(&self) -> Keypair {
+        Keypair::from_bytes(&self.0).unwrap()
+    }
+
+    pub fn pubkey(&self) -> Pubkey {
+        solana_sdk::signature::Signer::pubkey(&self.to_keypair())
+    }
+}
+impl Default for TestKeypair {
+    fn default() -> Self {
+        Self([0; 64])
+    }
+}
+impl<T: std::borrow::Borrow<Keypair>> From<T> for TestKeypair {
+    fn from(k: T) -> Self {
+        Self(k.borrow().to_bytes())
+    }
+}
+impl Into<Keypair> for &TestKeypair {
+    fn into(self) -> Keypair {
+        self.to_keypair()
+    }
+}

programs/mango-v3-reimbursement/tests/test_basic.rs

@@ -0,0 +1,415 @@
+#![cfg(feature = "test-bpf")]
+
+use solana_program_test::{BanksClientError, *};
+
+use mango_v3_reimbursement::state::*;
+use program_test::*;
+use std::sync::Arc;
+
+mod program_test;
+
+fn make_table(authority: Pubkey, rows: &[Row]) -> solana_sdk::account::Account {
+    let mut data = vec![0u8; 40 + rows.len() * 160];
+    data[5..37].copy_from_slice(&authority.to_bytes());
+    for (i, row) in rows.iter().enumerate() {
+        data[40 + i * 160..40 + (i + 1) * 160].copy_from_slice(bytemuck::bytes_of(row));
+    }
+    let mut acc =
+        solana_sdk::account::Account::new(u32::MAX as u64, data.len(), &Pubkey::new_unique());
+    acc.data.copy_from_slice(&data);
+    acc
+}
+
+async fn token_transfer(
+    solana: &Arc<SolanaCookie>,
+    from: Pubkey,
+    to: Pubkey,
+    authority: TestKeypair,
+    amount: u64,
+) -> std::result::Result<(), BanksClientError> {
+    let mut tx = ClientTransaction::new(solana);
+    tx.add_instruction_direct(
+        spl_token::instruction::transfer(
+            &spl_token::ID,
+            &from,
+            &to,
+            &authority.pubkey(),
+            &[&authority.pubkey()],
+            amount,
+        )
+        .unwrap(),
+    );
+    tx.add_signer(authority);
+    tx.send().await
+}
+
+async fn create_ata(
+    solana: &Arc<SolanaCookie>,
+    authority: Pubkey,
+    payer: TestKeypair,
+    mint: Pubkey,
+) -> std::result::Result<Pubkey, BanksClientError> {
+    let mut tx = ClientTransaction::new(solana);
+    tx.add_instruction_direct(
+        spl_associated_token_account::instruction::create_associated_token_account(
+            &payer.pubkey(),
+            &authority,
+            &mint,
+        ),
+    );
+    tx.add_signer(payer);
+    tx.send().await?;
+    Ok(spl_associated_token_account::get_associated_token_address(
+        &authority, &mint,
+    ))
+}
+
+async fn load_reimbursement(
+    solana: &SolanaCookie,
+    reimbursement_account: Pubkey,
+) -> ([bool; 16], [bool; 16]) {
+    let data: ReimbursementAccount = solana.get_account(reimbursement_account).await;
+    let reimb = (0..16).map(|i| data.reimbursed(i)).collect::<Vec<bool>>();
+    let transf = (0..16)
+        .map(|i| data.claim_transferred(i))
+        .collect::<Vec<bool>>();
+    (reimb.try_into().unwrap(), transf.try_into().unwrap())
+}
+
+#[tokio::test]
+async fn test_basic() -> Result<()> {
+    use mango_v3_reimbursement::accounts;
+
+    let authority = TestKeypair::new();
+    let table = Pubkey::new_unique();
+    let user1 = TestKeypair::new();
+    let user2 = TestKeypair::new();
+    let table_rows = vec![
+        Row {
+            owner: user1.pubkey(),
+            balances: (0..16)
+                .map(|i| 100 + i)
+                .collect::<Vec<u64>>()
+                .try_into()
+                .unwrap(),
+        },
+        Row {
+            owner: user2.pubkey(),
+            balances: (0..16).collect::<Vec<u64>>().try_into().unwrap(),
+        },
+    ];
+    let table_account = make_table(authority.pubkey(), &table_rows);
+
+    let mut builder = TestContextBuilder::new();
+    builder.test().add_account(table, table_account);
+
+    let context = builder.start_default().await;
+    let solana = &context.solana.clone();
+
+    let payer = context.users[1].key;
+
+    let mut user1_token = vec![];
+    for i in 0..2 {
+        user1_token.push(
+            create_ata(solana, user1.pubkey(), payer, context.mints[i].pubkey)
+                .await
+                .unwrap(),
+        );
+    }
+    let mut user2_token = vec![];
+    for i in 0..2 {
+        user2_token.push(
+            create_ata(solana, user2.pubkey(), payer, context.mints[i].pubkey)
+                .await
+                .unwrap(),
+        );
+    }
+
+    let accounts::CreateGroup { group, .. } = send_tx(
+        solana,
+        CreateGroupInstruction {
+            group_num: 0,
+            testing: false,
+            claim_transfer_destination: authority.pubkey(),
+            authority,
+            payer,
+            table,
+        },
+    )
+    .await
+    .unwrap();
+
+    let mut claim_accounts = vec![];
+    for i in [0, 1, 15] {
+        let accounts::CreateVault {
+            vault,
+            claim_transfer_token_account,
+            ..
+        } = send_tx(
+            solana,
+            CreateVaultInstruction {
+                group,
+                authority,
+                payer,
+                token_index: i,
+                mint: context.mints[i].pubkey,
+            },
+        )
+        .await
+        .unwrap();
+        token_transfer(
+            solana,
+            context.users[1].token_accounts[i],
+            vault,
+            payer,
+            1000,
+        )
+        .await
+        .unwrap();
+        claim_accounts.push(claim_transfer_token_account);
+    }
+
+    //
+    // Test reimbursing user1
+    //
+
+    let accounts::CreateReimbursementAccount {
+        reimbursement_account,
+        ..
+    } = send_tx(
+        solana,
+        CreateReimbursementAccountInstruction {
+            group,
+            mango_account_owner: user1.pubkey(),
+            payer,
+        },
+    )
+    .await
+    .unwrap();
+
+    // Creating twice is ok
+    send_tx(
+        solana,
+        CreateReimbursementAccountInstruction {
+            group,
+            mango_account_owner: user1.pubkey(),
+            payer,
+        },
+    )
+    .await
+    .unwrap();
+
+    // Cannot reimburse before start
+    assert!(send_tx(
+        solana,
+        ReimburseInstruction {
+            group,
+            token_index: 0,
+            table_index: 0,
+            mango_account_owner: user1,
+            transfer_claim: true,
+            token_account: user1_token[0],
+        }
+    )
+    .await
+    .is_err());
+
+    send_tx(solana, StartReimbursementInstruction { group, authority })
+        .await
+        .unwrap();
+
+    // Cannot reimburse a bad table index
+    assert!(send_tx(
+        solana,
+        ReimburseInstruction {
+            group,
+            token_index: 0,
+            table_index: 1,
+            mango_account_owner: user1,
+            transfer_claim: true,
+            token_account: user1_token[0],
+        }
+    )
+    .await
+    .is_err());
+
+    send_tx(
+        solana,
+        ReimburseInstruction {
+            group,
+            token_index: 0,
+            table_index: 0,
+            mango_account_owner: user1,
+            transfer_claim: false,
+            token_account: user1_token[0],
+        },
+    )
+    .await
+    .unwrap();
+
+    assert_eq!(solana.token_account_balance(user1_token[0]).await, 100);
+    assert_eq!(solana.token_account_balance(claim_accounts[0]).await, 0);
+    {
+        let (reimb, transf) = load_reimbursement(solana, reimbursement_account).await;
+        assert_eq!(reimb[0], true);
+        assert_eq!(reimb[1..], [false; 15]);
+        assert_eq!(transf, [false; 16]);
+    }
+
+    // does not work again (regardless of transfer_claim)
+    assert!(send_tx(
+        solana,
+        ReimburseInstruction {
+            group,
+            token_index: 0,
+            table_index: 0,
+            mango_account_owner: user1,
+            transfer_claim: true,
+            token_account: user1_token[0],
+        }
+    )
+    .await
+    .is_err());
+    assert!(send_tx(
+        solana,
+        ReimburseInstruction {
+            group,
+            token_index: 0,
+            table_index: 0,
+            mango_account_owner: user1,
+            transfer_claim: false,
+            token_account: user1_token[0],
+        }
+    )
+    .await
+    .is_err());
+
+    // can claim a second token
+    send_tx(
+        solana,
+        ReimburseInstruction {
+            group,
+            token_index: 1,
+            table_index: 0,
+            mango_account_owner: user1,
+            transfer_claim: true,
+            token_account: user1_token[1],
+        },
+    )
+    .await
+    .unwrap();
+
+    assert_eq!(solana.token_account_balance(user1_token[1]).await, 101);
+    assert_eq!(solana.token_account_balance(claim_accounts[1]).await, 101);
+    {
+        let (reimb, transf) = load_reimbursement(solana, reimbursement_account).await;
+        assert_eq!(reimb[0..2], [true; 2]);
+        assert_eq!(reimb[2..], [false; 14]);
+        assert_eq!(transf[0..2], [false, true]);
+        assert_eq!(transf[2..], [false; 14]);
+    }
+
+    // can't claim again
+    assert!(send_tx(
+        solana,
+        ReimburseInstruction {
+            group,
+            token_index: 1,
+            table_index: 0,
+            mango_account_owner: user1,
+            transfer_claim: false,
+            token_account: user1_token[1],
+        }
+    )
+    .await
+    .is_err());
+
+    //
+    // Test reimbursing user2
+    //
+
+    let accounts::CreateReimbursementAccount {
+        reimbursement_account,
+        ..
+    } = send_tx(
+        solana,
+        CreateReimbursementAccountInstruction {
+            group,
+            mango_account_owner: user2.pubkey(),
+            payer,
+        },
+    )
+    .await
+    .unwrap();
+
+    send_tx(
+        solana,
+        ReimburseInstruction {
+            group,
+            token_index: 1,
+            table_index: 1,
+            mango_account_owner: user2,
+            transfer_claim: true,
+            token_account: user2_token[1],
+        },
+    )
+    .await
+    .unwrap();
+
+    // ok to call reimburse, even if amount=0
+    send_tx(
+        solana,
+        ReimburseInstruction {
+            group,
+            token_index: 0,
+            table_index: 1,
+            mango_account_owner: user2,
+            transfer_claim: true,
+            token_account: user2_token[0],
+        },
+    )
+    .await
+    .unwrap();
+
+    assert_eq!(solana.token_account_balance(user2_token[0]).await, 0);
+    assert_eq!(solana.token_account_balance(user2_token[1]).await, 1);
+    assert_eq!(solana.token_account_balance(claim_accounts[1]).await, 101 + 1);
+    {
+        let (reimb, transf) = load_reimbursement(solana, reimbursement_account).await;
+        assert_eq!(reimb[0..2], [true; 2]);
+        assert_eq!(reimb[2..], [false; 14]);
+        assert_eq!(transf[0..2], [true; 2]);
+        assert_eq!(transf[2..], [false; 14]);
+    }
+
+    // Can claim the rightmost row entry too
+    let user2_token15 = create_ata(solana, user2.pubkey(), payer, context.mints[15].pubkey)
+        .await
+        .unwrap();
+    send_tx(
+        solana,
+        ReimburseInstruction {
+            group,
+            token_index: 15,
+            table_index: 1,
+            mango_account_owner: user2,
+            transfer_claim: false,
+            token_account: user2_token15,
+        },
+    )
+    .await
+    .unwrap();
+
+    assert_eq!(solana.token_account_balance(user2_token15).await, 15);
+    assert_eq!(solana.token_account_balance(claim_accounts[2]).await, 0);
+    {
+        let (reimb, transf) = load_reimbursement(solana, reimbursement_account).await;
+        assert_eq!(reimb[0..2], [true; 2]);
+        assert_eq!(reimb[2..15], [false; 13]);
+        assert_eq!(reimb[15], true);
+        assert_eq!(transf[0..2], [true; 2]);
+        assert_eq!(transf[2..], [false; 14]);
+    }
+
+    Ok(())
+}