From 3f9d235b9840b207da7afca6d5a7cf2c1eebaaa8 Mon Sep 17 00:00:00 2001 From: swimricky Date: Tue, 1 Aug 2023 16:43:56 -0400 Subject: [PATCH 1/4] feat: use external token account as treasury and approve config as delegate --- .../programs/token-dispenser/src/lib.rs | 37 ++-- .../src/tests/dispenser_simulator.rs | 130 +++++++++---- .../src/tests/test_checkout.rs | 182 +++++++++++++++++- .../src/tests/test_happy_path.rs | 31 ++- .../src/tests/test_initialize.rs | 48 +---- 5 files changed, 317 insertions(+), 111 deletions(-) diff --git a/token-dispenser/programs/token-dispenser/src/lib.rs b/token-dispenser/programs/token-dispenser/src/lib.rs index c9443b4c..35adc227 100644 --- a/token-dispenser/programs/token-dispenser/src/lib.rs +++ b/token-dispenser/programs/token-dispenser/src/lib.rs @@ -173,21 +173,17 @@ pub mod token_dispenser { #[derive(Accounts)] pub struct Initialize<'info> { #[account(mut)] - pub payer: Signer<'info>, + pub payer: Signer<'info>, #[account(init, payer = payer, space = Config::LEN, seeds = [CONFIG_SEED], bump)] - pub config: Account<'info, Config>, + pub config: Account<'info, Config>, /// Mint of the treasury - pub mint: Account<'info, Mint>, - #[account( - init, - payer = payer, - associated_token::authority = config, - associated_token::mint = mint, - )] - pub treasury: Account<'info, TokenAccount>, - pub system_program: Program<'info, System>, - pub token_program: Program<'info, Token>, - pub associated_token_program: Program<'info, AssociatedToken>, + pub mint: Account<'info, Mint>, + /// Treasury token account. This is an externally owned token account and + /// the owner of this account will approve the config as a delegate using the + /// solana CLI command `spl-token approve ` + #[account( token::mint = mint )] + pub treasury: Account<'info, TokenAccount>, + pub system_program: Program<'info, System>, } #[derive(Accounts)] @@ -221,11 +217,7 @@ pub struct Checkout<'info> { /// Mint of the treasury & claimant_fund token account. /// Needed if the `claimant_fund` token account needs to be initialized pub mint: Account<'info, Mint>, - #[account( - mut, - associated_token::authority = config, - associated_token::mint = config.mint, - )] + #[account(mut)] pub treasury: Account<'info, TokenAccount>, #[account(mut, seeds = [CART_SEED, claimant.key.as_ref()], bump)] pub cart: Account<'info, Cart>, @@ -544,10 +536,6 @@ pub fn get_cart_pda(claimant: &Pubkey) -> (Pubkey, u8) { Pubkey::find_program_address(&[CART_SEED, claimant.as_ref()], &crate::id()) } -pub fn get_treasury_ata(config: &Pubkey, mint: &Pubkey) -> Pubkey { - get_associated_token_address(config, mint) -} - impl crate::accounts::Initialize { pub fn populate(payer: Pubkey, mint: Pubkey, treasury: Pubkey) -> Self { crate::accounts::Initialize { @@ -556,8 +544,6 @@ impl crate::accounts::Initialize { mint, treasury, system_program: system_program::System::id(), - token_program: anchor_spl::token::Token::id(), - associated_token_program: AssociatedToken::id(), } } } @@ -579,6 +565,7 @@ impl crate::accounts::Checkout { pub fn populate( claimant: Pubkey, mint: Pubkey, + treasury: Pubkey, cart_override: Option, claimant_fund_override: Option, ) -> Self { @@ -587,7 +574,7 @@ impl crate::accounts::Checkout { claimant, config, mint, - treasury: get_treasury_ata(&config, &mint), + treasury, cart: cart_override.unwrap_or_else(|| get_cart_pda(&claimant).0), claimant_fund: claimant_fund_override .unwrap_or_else(|| get_associated_token_address(&claimant, &mint)), diff --git a/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs b/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs index 205096c7..db74a603 100644 --- a/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs +++ b/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs @@ -2,10 +2,7 @@ use { super::test_happy_path::TestClaimCertificate, crate::{ accounts, - get_cart_pda, - get_config_pda, get_receipt_pda, - get_treasury_ata, instruction, tests::merkleize, ClaimInfo, @@ -31,22 +28,19 @@ use { InstructionData, ToAccountMetas, }, - anchor_spl::{ - associated_token::{ - get_associated_token_address, - get_associated_token_address_with_program_id, - AssociatedToken, - }, - token::{ - spl_token::instruction::{ + anchor_spl::token::{ + spl_token, + spl_token::{ + error::TokenError, + instruction::{ initialize_account3, - initialize_mint, + initialize_mint2, mint_to, }, - Mint, - Token, - TokenAccount, }, + Mint, + Token, + TokenAccount, }, pythnet_sdk::accumulators::merkle::{ MerkleRoot, @@ -74,10 +68,13 @@ use { }; pub struct DispenserSimulator { - banks_client: BanksClient, - pub genesis_keypair: Keypair, - recent_blockhash: hash::Hash, - pub mint_keypair: Keypair, + banks_client: BanksClient, + pub genesis_keypair: Keypair, + recent_blockhash: hash::Hash, + pub mint_keypair: Keypair, + /// also the owner/authority of `pyth_treasury` + pub pyth_mint_authority: Keypair, + pub pyth_treasury: Pubkey, } impl DispenserSimulator { @@ -85,20 +82,39 @@ impl DispenserSimulator { let program_test = ProgramTest::new("token_dispenser", crate::id(), None); let (banks_client, genesis_keypair, recent_blockhash) = program_test.start().await; let mint_keypair = Keypair::new(); + let pyth_mint_authority = Keypair::new(); + let pyth_treasury = Keypair::new(); let mut simulator = DispenserSimulator { banks_client, genesis_keypair, recent_blockhash, mint_keypair, + pyth_mint_authority, + pyth_treasury: pyth_treasury.pubkey(), }; + simulator .create_mint( ©_keypair(&simulator.mint_keypair), - ©_keypair(&simulator.genesis_keypair), + &simulator.pyth_mint_authority.pubkey(), 6, ) .await .unwrap(); + + simulator + .create_token_account( + simulator.mint_keypair.pubkey(), + ©_keypair(&simulator.pyth_mint_authority), + &pyth_treasury, + ) + .await + .unwrap(); + // simulator + // .create_associated_token_account( + // simulator.mint_keypair.pubkey(), + // ©_keypair(&simulator.pyth_mint_authority), + // ).await.unwrap(); simulator } @@ -118,7 +134,7 @@ impl DispenserSimulator { pub async fn create_mint( &mut self, mint_keypair: &Keypair, - mint_authority: &Keypair, + mint_authority: &Pubkey, decimals: u8, ) -> Result<(), BanksClientError> { let space = Mint::LEN; @@ -131,32 +147,29 @@ impl DispenserSimulator { space as u64, &Token::id(), ), - initialize_mint( + initialize_mint2( &Token::id(), &mint_keypair.pubkey(), - &mint_authority.pubkey(), + &mint_authority, None, decimals, ) .unwrap(), ]; - self.process_ix(init_mint_ixs, &vec![mint_keypair, mint_authority]) - .await + self.process_ix(init_mint_ixs, &vec![mint_keypair]).await } pub async fn mint_to_treasury(&mut self, mint_amount: u64) -> Result<(), BanksClientError> { - let config = get_config_pda().0; - let treasury = get_treasury_ata(&config, &self.mint_keypair.pubkey()); let mint_to_ix = &[mint_to( &Token::id(), &self.mint_keypair.pubkey(), - &treasury, - &self.genesis_keypair.pubkey(), + &self.pyth_treasury, + &self.pyth_mint_authority.pubkey(), &[], mint_amount, ) .unwrap()]; - self.process_ix(mint_to_ix, &vec![©_keypair(&self.genesis_keypair)]) + self.process_ix(mint_to_ix, &vec![©_keypair(&self.pyth_mint_authority)]) .await } @@ -187,15 +200,10 @@ impl DispenserSimulator { mint_pubkey_override: Option, treasury_pubkey_override: Option, ) -> Result<(), BanksClientError> { - let treasury = get_associated_token_address_with_program_id( - &(get_config_pda().0), - &self.mint_keypair.pubkey(), - &Token::id(), - ); let accounts = accounts::Initialize::populate( self.genesis_keypair.pubkey(), mint_pubkey_override.unwrap_or(self.mint_keypair.pubkey()), - treasury_pubkey_override.unwrap_or(treasury), + treasury_pubkey_override.unwrap_or(self.pyth_treasury), ) .to_account_metas(None); let instruction_data = instruction::Initialize { @@ -247,6 +255,44 @@ impl DispenserSimulator { Ok((merkle_tree, mock_offchain_certificates_and_claimants)) } + pub async fn approve_treasury_delegate( + &mut self, + delegate: Pubkey, + amount: u64, + ) -> Result<(), BanksClientError> { + let approve_ix = spl_token::instruction::approve( + &Token::id(), + &self.pyth_treasury, + &delegate, + &self.pyth_mint_authority.pubkey(), + &[], + amount, + ) + .unwrap(); + + self.process_ix( + &[approve_ix], + &vec![©_keypair(&self.pyth_mint_authority)], + ) + .await + } + + pub async fn revoke_treasury_delegate(&mut self) -> Result<(), BanksClientError> { + let revoke_ix = spl_token::instruction::revoke( + &Token::id(), + &self.pyth_treasury, + &self.pyth_mint_authority.pubkey(), + &[], + ) + .unwrap(); + + self.process_ix( + &[revoke_ix], + &vec![©_keypair(&self.pyth_mint_authority)], + ) + .await + } + pub async fn claim( &mut self, @@ -307,6 +353,7 @@ impl DispenserSimulator { let accounts = accounts::Checkout::populate( claimant.pubkey(), mint, + self.pyth_treasury, cart_override, claimant_fund_override, ) @@ -358,7 +405,7 @@ impl DispenserSimulator { ) .unwrap(), ]; - self.process_ix(init_token_account_ixs, &vec![&owner, &token_account]) + self.process_ix(init_token_account_ixs, &vec![&token_account]) .await } } @@ -386,6 +433,15 @@ impl IntoTransactionError for ErrorCode { ) } } + +impl IntoTransactionError for TokenError { + fn into_transaction_error(self) -> TransactionError { + TransactionError::InstructionError( + 0, + InstructionError::try_from(u64::from(ProgramError::from(self))).unwrap(), + ) + } +} impl IntoTransactionError for InstructionError { fn into_transaction_error(self) -> TransactionError { TransactionError::InstructionError(0, self) diff --git a/token-dispenser/programs/token-dispenser/src/tests/test_checkout.rs b/token-dispenser/programs/token-dispenser/src/tests/test_checkout.rs index 2b067a1a..e4ef9cc3 100644 --- a/token-dispenser/programs/token-dispenser/src/tests/test_checkout.rs +++ b/token-dispenser/programs/token-dispenser/src/tests/test_checkout.rs @@ -2,6 +2,7 @@ use { super::dispenser_simulator::DispenserSimulator, crate::{ get_cart_pda, + get_config_pda, tests::dispenser_simulator::{ copy_keypair, IntoTransactionError, @@ -10,11 +11,18 @@ use { }, anchor_lang::solana_program::{ instruction::InstructionError::MissingAccount, + program_option::COption, system_instruction, }, anchor_spl::{ associated_token::get_associated_token_address, - token::TokenAccount, + token::{ + spl_token::error::TokenError::{ + InsufficientFunds, + OwnerMismatch, + }, + TokenAccount, + }, }, solana_program_test::tokio, solana_sdk::{ @@ -195,6 +203,20 @@ pub async fn test_checkout_fails_with_insufficient_funds() { let total_claim_sum = claim_sums.iter().sum::(); simulator.mint_to_treasury(total_claim_sum).await.unwrap(); + // approve enough for first checkout + let delegated_amount = claim_sums[0] + 1; + simulator + .approve_treasury_delegate(get_config_pda().0, delegated_amount) + .await + .unwrap(); + + let treasury_data = simulator + .get_account_data::(simulator.pyth_treasury) + .await + .unwrap(); + assert_eq!(treasury_data.amount, total_claim_sum); + assert_eq!(treasury_data.delegated_amount, delegated_amount); + simulator .checkout( ©_keypair(&simulator.genesis_keypair), @@ -204,6 +226,64 @@ pub async fn test_checkout_fails_with_insufficient_funds() { ) .await .unwrap(); + let treasury_data = simulator + .get_account_data::(simulator.pyth_treasury) + .await + .unwrap(); + assert_eq!(treasury_data.amount, total_claim_sum - claim_sums[0]); + assert_eq!(treasury_data.delegated_amount, 1); + assert_eq!( + simulator + .checkout( + ©_keypair(&claimant_1), + simulator.mint_keypair.pubkey(), + None, + None, + ) + .await + .unwrap_err() + .unwrap(), + InsufficientFunds.into_transaction_error() + ); + + let delegated_amount = claim_sums[1] - 1; + simulator + .approve_treasury_delegate(get_config_pda().0, delegated_amount) + .await + .unwrap(); + let treasury_data = simulator + .get_account_data::(simulator.pyth_treasury) + .await + .unwrap(); + assert_eq!(treasury_data.amount, total_claim_sum - claim_sums[0]); + assert_eq!(treasury_data.delegated_amount, delegated_amount); + + assert_eq!( + simulator + .checkout( + ©_keypair(&claimant_1), + simulator.mint_keypair.pubkey(), + None, + None, + ) + .await + .unwrap_err() + .unwrap(), + InsufficientFunds.into_transaction_error() + ); + + let delegated_amount = claim_sums[1]; + simulator + .approve_treasury_delegate(get_config_pda().0, delegated_amount) + .await + .unwrap(); + + let treasury_data = simulator + .get_account_data::(simulator.pyth_treasury) + .await + .unwrap(); + assert_eq!(treasury_data.delegated_amount, delegated_amount); + assert_eq!(treasury_data.delegate, COption::Some(get_config_pda().0)); simulator .checkout( @@ -215,6 +295,16 @@ pub async fn test_checkout_fails_with_insufficient_funds() { .await .unwrap(); + let treasury_data = simulator + .get_account_data::(simulator.pyth_treasury) + .await + .unwrap(); + + assert_eq!(treasury_data.amount, 0); + assert_eq!(treasury_data.delegated_amount, 0); + assert_eq!(treasury_data.delegate, COption::None); + + let claimant_pubkeys = vec![simulator.genesis_keypair.pubkey(), claimant_1.pubkey()]; for (claim_sum, pubkey) in claim_sums.iter().zip(claimant_pubkeys.iter()) { let claimant_fund_data = simulator @@ -227,3 +317,93 @@ pub async fn test_checkout_fails_with_insufficient_funds() { assert_eq!(claimant_fund_data.amount, *claim_sum); } } + +#[tokio::test] +pub async fn test_checkout_fails_if_delegate_revoked() { + let dispenser_guard: Keypair = Keypair::new(); + + let mut simulator = DispenserSimulator::new().await; + let claimant_1 = Keypair::new(); + let claimant_1_airdrop_ix = system_instruction::transfer( + &simulator.genesis_keypair.pubkey(), + &claimant_1.pubkey(), + 1000000000, + ); + + simulator + .process_ix(&vec![claimant_1_airdrop_ix], &vec![]) + .await + .unwrap(); + + + let (merkle_tree, mock_offchain_certificates_and_claimants) = simulator + .initialize_with_claimants( + vec![ + copy_keypair(&simulator.genesis_keypair), + copy_keypair(&claimant_1), + ], + &dispenser_guard, + ) + .await + .unwrap(); + + for (claimant, offchain_claim_certificates) in &mock_offchain_certificates_and_claimants { + for offchain_claim_certificate in offchain_claim_certificates { + simulator + .claim( + ©_keypair(claimant), + &dispenser_guard, + &offchain_claim_certificate, + &merkle_tree, + ) + .await + .unwrap(); + } + } + + + let claim_sums = mock_offchain_certificates_and_claimants + .iter() + .map(|x| x.1.iter().map(|y| y.amount).sum::()) + .collect::>(); + + let total_claim_sum = claim_sums.iter().sum::(); + simulator.mint_to_treasury(total_claim_sum).await.unwrap(); + + + simulator + .approve_treasury_delegate(get_config_pda().0, total_claim_sum) + .await + .unwrap(); + + let treasury_data = simulator + .get_account_data::(simulator.pyth_treasury) + .await + .unwrap(); + assert_eq!(treasury_data.amount, total_claim_sum); + assert_eq!(treasury_data.delegated_amount, total_claim_sum); + assert_eq!(treasury_data.delegate, COption::Some(get_config_pda().0)); + + simulator.revoke_treasury_delegate().await.unwrap(); + let treasury_data = simulator + .get_account_data::(simulator.pyth_treasury) + .await + .unwrap(); + assert_eq!(treasury_data.amount, total_claim_sum); + assert_eq!(treasury_data.delegated_amount, 0); + assert_eq!(treasury_data.delegate, COption::None); + + assert_eq!( + simulator + .checkout( + ©_keypair(&claimant_1), + simulator.mint_keypair.pubkey(), + None, + None, + ) + .await + .unwrap_err() + .unwrap(), + OwnerMismatch.into_transaction_error() + ); +} diff --git a/token-dispenser/programs/token-dispenser/src/tests/test_happy_path.rs b/token-dispenser/programs/token-dispenser/src/tests/test_happy_path.rs index ac18a651..f06ecc2f 100644 --- a/token-dispenser/programs/token-dispenser/src/tests/test_happy_path.rs +++ b/token-dispenser/programs/token-dispenser/src/tests/test_happy_path.rs @@ -7,7 +7,6 @@ use { get_cart_pda, get_config_pda, get_receipt_pda, - get_treasury_ata, tests::{ dispenser_simulator::{ copy_keypair, @@ -26,7 +25,7 @@ use { }, anchor_lang::{ prelude::Pubkey, - AccountDeserialize, + solana_program::program_option::COption, AnchorDeserialize, AnchorSerialize, }, @@ -184,7 +183,7 @@ pub async fn test_happy_path() { let (merkle_tree, merkle_items_serialized) = merkleize(merkle_items); let (config_pubkey, config_bump) = get_config_pda(); - let treasury = get_treasury_ata(&config_pubkey, &simulator.mint_keypair.pubkey()); + let treasury = simulator.pyth_treasury; simulator .initialize( @@ -213,14 +212,29 @@ pub async fn test_happy_path() { .map(|item| item.amount) .sum::(); let mint_to_amount = 10 * claim_sum; - simulator.mint_to_treasury(mint_to_amount).await.unwrap(); - let treasury_account: Account = simulator.get_account(treasury).await.unwrap(); - let treasury_data: TokenAccount = - TokenAccount::try_deserialize_unchecked(&mut treasury_account.data.as_slice()).unwrap(); + simulator.mint_to_treasury(mint_to_amount).await.unwrap(); + let treasury_data = simulator + .get_account_data::(treasury) + .await + .unwrap(); assert_eq!(treasury_data.amount, mint_to_amount); assert_eq!(treasury_data.mint, simulator.mint_keypair.pubkey()); - assert_eq!(treasury_data.owner, config_pubkey); + assert_eq!(treasury_data.owner, simulator.pyth_mint_authority.pubkey()); + assert_eq!(treasury_data.delegate, COption::None); + assert_eq!(treasury_data.delegated_amount, 0); + + simulator + .approve_treasury_delegate(get_config_pda().0, mint_to_amount) + .await + .unwrap(); + + let treasury_data = simulator + .get_account_data::(treasury) + .await + .unwrap(); + assert_eq!(treasury_data.delegate, COption::Some(config_pubkey)); + assert_eq!(treasury_data.delegated_amount, mint_to_amount); for serialized_item in &merkle_items_serialized { assert!(simulator @@ -270,6 +284,7 @@ pub async fn test_happy_path() { .unwrap(); assert_eq!(cart_data.amount, claim_sum); + // Checkout simulator .checkout( diff --git a/token-dispenser/programs/token-dispenser/src/tests/test_initialize.rs b/token-dispenser/programs/token-dispenser/src/tests/test_initialize.rs index f517d1e6..4acf10b2 100644 --- a/token-dispenser/programs/token-dispenser/src/tests/test_initialize.rs +++ b/token-dispenser/programs/token-dispenser/src/tests/test_initialize.rs @@ -1,7 +1,6 @@ use { super::dispenser_simulator::DispenserSimulator, crate::{ - get_config_pda, tests::{ dispenser_simulator::copy_keypair, merkleize, @@ -9,11 +8,7 @@ use { }, ClaimInfo, }, - anchor_lang::solana_program::instruction::InstructionError::{ - Custom, - MissingAccount, - }, - anchor_spl::associated_token::get_associated_token_address, + anchor_lang::solana_program::instruction::InstructionError::Custom, solana_program_test::tokio, solana_sdk::{ signature::Keypair, @@ -43,8 +38,6 @@ pub async fn test_initialize_fails_with_incorrect_accounts() { let (merkle_tree, _) = merkleize(merkle_items); - let config_pubkey = get_config_pda().0; - let res = simulator .initialize( @@ -61,38 +54,16 @@ pub async fn test_initialize_fails_with_incorrect_accounts() { let fake_mint_keypair = Keypair::new(); simulator - .create_mint( - &fake_mint_keypair, - ©_keypair(&simulator.genesis_keypair), - 0, - ) + .create_mint(&fake_mint_keypair, &simulator.genesis_keypair.pubkey(), 0) .await .unwrap(); // create treasury (associated token account) from a different mint - let invalid_treasury = - get_associated_token_address(&(config_pubkey), &fake_mint_keypair.pubkey()); - - let res = simulator - .initialize( - merkle_tree.root.clone(), - dispenser_guard.pubkey(), - None, - Some(invalid_treasury), - ) - .await; - assert!(res.is_err()); - assert_eq!( - res.unwrap_err().unwrap(), - InstructionError(0, MissingAccount) - ); - - // create token account with correct mint but not an associated token account - let fake_treasury_keypair = Keypair::new(); + let invalid_treasury = Keypair::new(); simulator .create_token_account( - simulator.mint_keypair.pubkey(), + fake_mint_keypair.pubkey(), ©_keypair(&simulator.genesis_keypair), - &fake_treasury_keypair, + &invalid_treasury, ) .await .unwrap(); @@ -102,13 +73,10 @@ pub async fn test_initialize_fails_with_incorrect_accounts() { merkle_tree.root.clone(), dispenser_guard.pubkey(), None, - Some(fake_treasury_keypair.pubkey()), //incorrect treasury + Some(invalid_treasury.pubkey()), ) .await; assert!(res.is_err()); - // associated token account for treasury is missing - assert_eq!( - res.unwrap_err().unwrap(), - InstructionError(0, MissingAccount) - ); + // 2014- ConstraintTokenMint + assert_eq!(res.unwrap_err().unwrap(), InstructionError(0, Custom(2014))); } From 442ef82579ad2e2e419448d43a8707e61d82d905 Mon Sep 17 00:00:00 2001 From: swimricky Date: Tue, 1 Aug 2023 17:03:44 -0400 Subject: [PATCH 2/4] refactor: add test utility method --- .../src/tests/dispenser_simulator.rs | 21 +++ .../src/tests/test_checkout.rs | 144 +++++++++++++----- .../src/tests/test_happy_path.rs | 43 +++--- 3 files changed, 146 insertions(+), 62 deletions(-) diff --git a/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs b/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs index db74a603..0920a7d1 100644 --- a/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs +++ b/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs @@ -19,6 +19,7 @@ use { solana_program::{ hash, instruction::Instruction, + program_option::COption, system_instruction::create_account, }, system_program, @@ -408,6 +409,26 @@ impl DispenserSimulator { self.process_ix(init_token_account_ixs, &vec![&token_account]) .await } + + pub async fn verify_token_account_data( + &mut self, + token_account: Pubkey, + expected_amount: u64, + expected_delegate: COption, + expected_delegated_amount: u64, + ) -> Result<(), BanksClientError> { + let token_account_data = self + .get_account_data::(token_account) + .await + .unwrap(); + assert_eq!(token_account_data.amount, expected_amount); + assert_eq!(token_account_data.delegate, expected_delegate); + assert_eq!( + token_account_data.delegated_amount, + expected_delegated_amount + ); + Ok(()) + } } pub fn copy_keypair(keypair: &Keypair) -> Keypair { diff --git a/token-dispenser/programs/token-dispenser/src/tests/test_checkout.rs b/token-dispenser/programs/token-dispenser/src/tests/test_checkout.rs index e4ef9cc3..897e5775 100644 --- a/token-dispenser/programs/token-dispenser/src/tests/test_checkout.rs +++ b/token-dispenser/programs/token-dispenser/src/tests/test_checkout.rs @@ -116,7 +116,6 @@ pub async fn test_checkout_fails_with_wrong_accounts() { Some(fake_claimant_fund_keypair.pubkey()), ) .await; - // simulator.process_ix(&vec![checkout_ix], &vec![]).await; assert!(res.is_err()); // 3014 - AccountNotAssociatedTokenAccount @@ -210,12 +209,16 @@ pub async fn test_checkout_fails_with_insufficient_funds() { .await .unwrap(); - let treasury_data = simulator - .get_account_data::(simulator.pyth_treasury) + simulator + .verify_token_account_data( + simulator.pyth_treasury, + total_claim_sum, + COption::Some(get_config_pda().0), + delegated_amount, + ) .await .unwrap(); - assert_eq!(treasury_data.amount, total_claim_sum); - assert_eq!(treasury_data.delegated_amount, delegated_amount); + simulator .checkout( @@ -226,12 +229,17 @@ pub async fn test_checkout_fails_with_insufficient_funds() { ) .await .unwrap(); - let treasury_data = simulator - .get_account_data::(simulator.pyth_treasury) + + simulator + .verify_token_account_data( + simulator.pyth_treasury, + total_claim_sum - claim_sums[0], + COption::Some(get_config_pda().0), + 1, + ) .await .unwrap(); - assert_eq!(treasury_data.amount, total_claim_sum - claim_sums[0]); - assert_eq!(treasury_data.delegated_amount, 1); + assert_eq!( simulator .checkout( @@ -251,12 +259,16 @@ pub async fn test_checkout_fails_with_insufficient_funds() { .approve_treasury_delegate(get_config_pda().0, delegated_amount) .await .unwrap(); - let treasury_data = simulator - .get_account_data::(simulator.pyth_treasury) + + simulator + .verify_token_account_data( + simulator.pyth_treasury, + total_claim_sum - claim_sums[0], + COption::Some(get_config_pda().0), + delegated_amount, + ) .await .unwrap(); - assert_eq!(treasury_data.amount, total_claim_sum - claim_sums[0]); - assert_eq!(treasury_data.delegated_amount, delegated_amount); assert_eq!( simulator @@ -278,12 +290,16 @@ pub async fn test_checkout_fails_with_insufficient_funds() { .await .unwrap(); - let treasury_data = simulator - .get_account_data::(simulator.pyth_treasury) + simulator + .verify_token_account_data( + simulator.pyth_treasury, + total_claim_sum - claim_sums[0], + COption::Some(get_config_pda().0), + delegated_amount, + ) .await .unwrap(); - assert_eq!(treasury_data.delegated_amount, delegated_amount); - assert_eq!(treasury_data.delegate, COption::Some(get_config_pda().0)); + simulator .checkout( @@ -295,26 +311,24 @@ pub async fn test_checkout_fails_with_insufficient_funds() { .await .unwrap(); - let treasury_data = simulator - .get_account_data::(simulator.pyth_treasury) + + simulator + .verify_token_account_data(simulator.pyth_treasury, 0, COption::None, 0) .await .unwrap(); - assert_eq!(treasury_data.amount, 0); - assert_eq!(treasury_data.delegated_amount, 0); - assert_eq!(treasury_data.delegate, COption::None); - let claimant_pubkeys = vec![simulator.genesis_keypair.pubkey(), claimant_1.pubkey()]; for (claim_sum, pubkey) in claim_sums.iter().zip(claimant_pubkeys.iter()) { - let claimant_fund_data = simulator - .get_account_data::(get_associated_token_address( - &pubkey, - &simulator.mint_keypair.pubkey(), - )) + simulator + .verify_token_account_data( + get_associated_token_address(&pubkey, &simulator.mint_keypair.pubkey()), + *claim_sum, + COption::None, + 0, + ) .await .unwrap(); - assert_eq!(claimant_fund_data.amount, *claim_sum); } } @@ -376,27 +390,27 @@ pub async fn test_checkout_fails_if_delegate_revoked() { .await .unwrap(); - let treasury_data = simulator - .get_account_data::(simulator.pyth_treasury) + simulator + .verify_token_account_data( + simulator.pyth_treasury, + total_claim_sum, + COption::Some(get_config_pda().0), + total_claim_sum, + ) .await .unwrap(); - assert_eq!(treasury_data.amount, total_claim_sum); - assert_eq!(treasury_data.delegated_amount, total_claim_sum); - assert_eq!(treasury_data.delegate, COption::Some(get_config_pda().0)); + simulator.revoke_treasury_delegate().await.unwrap(); - let treasury_data = simulator - .get_account_data::(simulator.pyth_treasury) + simulator + .verify_token_account_data(simulator.pyth_treasury, total_claim_sum, COption::None, 0) .await .unwrap(); - assert_eq!(treasury_data.amount, total_claim_sum); - assert_eq!(treasury_data.delegated_amount, 0); - assert_eq!(treasury_data.delegate, COption::None); assert_eq!( simulator .checkout( - ©_keypair(&claimant_1), + ©_keypair(&simulator.genesis_keypair), simulator.mint_keypair.pubkey(), None, None, @@ -406,4 +420,54 @@ pub async fn test_checkout_fails_if_delegate_revoked() { .unwrap(), OwnerMismatch.into_transaction_error() ); + + simulator + .approve_treasury_delegate(get_config_pda().0, claim_sums[0]) + .await + .unwrap(); + + simulator + .verify_token_account_data( + simulator.pyth_treasury, + total_claim_sum, + COption::Some(get_config_pda().0), + claim_sums[0], + ) + .await + .unwrap(); + + simulator + .checkout( + ©_keypair(&simulator.genesis_keypair), + simulator.mint_keypair.pubkey(), + None, + None, + ) + .await + .unwrap(); + + + simulator + .verify_token_account_data( + simulator.pyth_treasury, + total_claim_sum - claim_sums[0], + COption::None, + 0, + ) + .await + .unwrap(); + + + simulator + .verify_token_account_data( + get_associated_token_address( + ©_keypair(&simulator.genesis_keypair).pubkey(), + &simulator.mint_keypair.pubkey(), + ), + claim_sums[0], + COption::None, + 0, + ) + .await + .unwrap(); } diff --git a/token-dispenser/programs/token-dispenser/src/tests/test_happy_path.rs b/token-dispenser/programs/token-dispenser/src/tests/test_happy_path.rs index f06ecc2f..220a9abb 100644 --- a/token-dispenser/programs/token-dispenser/src/tests/test_happy_path.rs +++ b/token-dispenser/programs/token-dispenser/src/tests/test_happy_path.rs @@ -29,10 +29,7 @@ use { AnchorDeserialize, AnchorSerialize, }, - anchor_spl::{ - associated_token::get_associated_token_address, - token::TokenAccount, - }, + anchor_spl::associated_token::get_associated_token_address, pythnet_sdk::accumulators::{ merkle::MerkleTree, Accumulator, @@ -214,27 +211,25 @@ pub async fn test_happy_path() { let mint_to_amount = 10 * claim_sum; simulator.mint_to_treasury(mint_to_amount).await.unwrap(); - let treasury_data = simulator - .get_account_data::(treasury) + simulator + .verify_token_account_data(treasury, mint_to_amount, COption::None, 0) .await .unwrap(); - assert_eq!(treasury_data.amount, mint_to_amount); - assert_eq!(treasury_data.mint, simulator.mint_keypair.pubkey()); - assert_eq!(treasury_data.owner, simulator.pyth_mint_authority.pubkey()); - assert_eq!(treasury_data.delegate, COption::None); - assert_eq!(treasury_data.delegated_amount, 0); simulator .approve_treasury_delegate(get_config_pda().0, mint_to_amount) .await .unwrap(); - let treasury_data = simulator - .get_account_data::(treasury) + simulator + .verify_token_account_data( + treasury, + mint_to_amount, + COption::Some(config_pubkey), + mint_to_amount, + ) .await .unwrap(); - assert_eq!(treasury_data.delegate, COption::Some(config_pubkey)); - assert_eq!(treasury_data.delegated_amount, mint_to_amount); for serialized_item in &merkle_items_serialized { assert!(simulator @@ -296,16 +291,20 @@ pub async fn test_happy_path() { .await .unwrap(); - let claimant_fund_data = simulator - .get_account_data::(get_associated_token_address( - &simulator.genesis_keypair.pubkey(), - &simulator.mint_keypair.pubkey(), - )) + + simulator + .verify_token_account_data( + get_associated_token_address( + &simulator.genesis_keypair.pubkey(), + &simulator.mint_keypair.pubkey(), + ), + claim_sum, + COption::None, + 0, + ) .await .unwrap(); - assert_eq!(claimant_fund_data.amount, claim_sum); - let cart_data = simulator .get_account_data::(cart_pda) .await From 38bdd179eaca8fa7a9b8b653eaf275233e4a6523 Mon Sep 17 00:00:00 2001 From: swimricky Date: Tue, 1 Aug 2023 17:08:06 -0400 Subject: [PATCH 3/4] style: delete commented out code --- .../token-dispenser/src/tests/dispenser_simulator.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs b/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs index 0920a7d1..046b32d9 100644 --- a/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs +++ b/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs @@ -111,11 +111,7 @@ impl DispenserSimulator { ) .await .unwrap(); - // simulator - // .create_associated_token_account( - // simulator.mint_keypair.pubkey(), - // ©_keypair(&simulator.pyth_mint_authority), - // ).await.unwrap(); + simulator } From daba443567999b427ccc44476b676d35ceaa422c Mon Sep 17 00:00:00 2001 From: swimricky Date: Wed, 2 Aug 2023 10:24:04 -0400 Subject: [PATCH 4/4] refactor: add airdrop util method in DispenserSimulator --- .../src/tests/dispenser_simulator.rs | 11 ++++++++- .../src/tests/test_checkout.rs | 24 +++---------------- 2 files changed, 13 insertions(+), 22 deletions(-) diff --git a/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs b/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs index 046b32d9..674641ba 100644 --- a/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs +++ b/token-dispenser/programs/token-dispenser/src/tests/dispenser_simulator.rs @@ -20,7 +20,10 @@ use { hash, instruction::Instruction, program_option::COption, - system_instruction::create_account, + system_instruction::{ + self, + create_account, + }, }, system_program, AccountDeserialize, @@ -128,6 +131,12 @@ impl DispenserSimulator { self.banks_client.get_rent().await.unwrap() } + pub async fn airdrop(&mut self, target: Pubkey, amount: u64) -> Result<(), BanksClientError> { + let airdrop_ix = + system_instruction::transfer(&self.genesis_keypair.pubkey(), &target, amount); + self.process_ix(&[airdrop_ix], &vec![]).await + } + pub async fn create_mint( &mut self, mint_keypair: &Keypair, diff --git a/token-dispenser/programs/token-dispenser/src/tests/test_checkout.rs b/token-dispenser/programs/token-dispenser/src/tests/test_checkout.rs index 897e5775..4a5cefe1 100644 --- a/token-dispenser/programs/token-dispenser/src/tests/test_checkout.rs +++ b/token-dispenser/programs/token-dispenser/src/tests/test_checkout.rs @@ -40,14 +40,8 @@ pub async fn test_checkout_fails_with_wrong_accounts() { let mut simulator = DispenserSimulator::new().await; let claimant_1 = Keypair::new(); - let claimant_1_airdrop_ix = system_instruction::transfer( - &simulator.genesis_keypair.pubkey(), - &claimant_1.pubkey(), - 1000000000, - ); - simulator - .process_ix(&vec![claimant_1_airdrop_ix], &vec![]) + .airdrop(claimant_1.pubkey(), 1000000000) .await .unwrap(); @@ -142,14 +136,8 @@ pub async fn test_checkout_fails_with_insufficient_funds() { let mut simulator = DispenserSimulator::new().await; let claimant_1 = Keypair::new(); - let claimant_1_airdrop_ix = system_instruction::transfer( - &simulator.genesis_keypair.pubkey(), - &claimant_1.pubkey(), - 1000000000, - ); - simulator - .process_ix(&vec![claimant_1_airdrop_ix], &vec![]) + .airdrop(claimant_1.pubkey(), 1000000000) .await .unwrap(); @@ -338,14 +326,8 @@ pub async fn test_checkout_fails_if_delegate_revoked() { let mut simulator = DispenserSimulator::new().await; let claimant_1 = Keypair::new(); - let claimant_1_airdrop_ix = system_instruction::transfer( - &simulator.genesis_keypair.pubkey(), - &claimant_1.pubkey(), - 1000000000, - ); - simulator - .process_ix(&vec![claimant_1_airdrop_ix], &vec![]) + .airdrop(claimant_1.pubkey(), 1000000000) .await .unwrap();