solana: Add Rust test cases to handle tokens with transfer fee (#489)

nvsriram · web-flow · commit a6fcaeaafd98 · 2024-08-02T16:37:23.000-04:00
* solana: Add sdk helper methods to use specified token program

* solana: Add test cases for locking and burning tokens with transfer fee
diff --git a/solana/programs/example-native-token-transfers/tests/common/setup.rs b/solana/programs/example-native-token-transfers/tests/common/setup.rs
@@ -15,14 +15,13 @@ use solana_sdk::{
     transaction::Transaction,
 };
 use spl_associated_token_account::get_associated_token_address_with_program_id;
-use spl_token::instruction::AuthorityType;
 use wormhole_anchor_sdk::wormhole::{BridgeData, FeeCollector};
 
 use crate::sdk::{
     accounts::{Governance, Wormhole, NTT},
     instructions::{
         admin::{register_transceiver, set_peer, RegisterTransceiver, SetPeer},
-        initialize::{initialize, Initialize},
+        initialize::{initialize_with_token_program_id, Initialize},
     },
     transceivers::wormhole::instructions::admin::{set_transceiver_peer, SetTransceiverPeer},
 };
@@ -73,10 +72,33 @@ pub async fn setup_with_extra_accounts(
     (ctx, test_data)
 }
 
+pub async fn setup_with_extra_accounts_with_transfer_fee(
+    mode: Mode,
+    accounts: &[(Pubkey, Account)],
+) -> (ProgramTestContext, TestData) {
+    let program_owner = Keypair::new();
+    let mut program_test = setup_programs(program_owner.pubkey()).await.unwrap();
+
+    for (pubkey, account) in accounts {
+        program_test.add_account(*pubkey, account.clone());
+    }
+
+    let mut ctx = program_test.start_with_context().await;
+
+    let test_data = setup_accounts_with_transfer_fee(&mut ctx, program_owner).await;
+    setup_ntt_with_token_program_id(&mut ctx, &test_data, mode, &spl_token_2022::id()).await;
+
+    (ctx, test_data)
+}
+
 pub async fn setup(mode: Mode) -> (ProgramTestContext, TestData) {
     setup_with_extra_accounts(mode, &[]).await
 }
 
+pub async fn setup_with_transfer_fee(mode: Mode) -> (ProgramTestContext, TestData) {
+    setup_with_extra_accounts_with_transfer_fee(mode, &[]).await
+}
+
 fn prefer_bpf() -> bool {
     std::env::var("BPF_OUT_DIR").is_ok() || std::env::var("SBF_OUT_DIR").is_ok()
 }
@@ -126,13 +148,22 @@ pub async fn setup_programs(program_owner: Pubkey) -> Result<ProgramTest, Error>
 /// Set up test accounts, and mint MINT_AMOUNT to the user's token account
 /// Set up the program for locking mode, and registers a peer
 pub async fn setup_ntt(ctx: &mut ProgramTestContext, test_data: &TestData, mode: Mode) {
+    setup_ntt_with_token_program_id(ctx, test_data, mode, &Token::id()).await;
+}
+
+pub async fn setup_ntt_with_token_program_id(
+    ctx: &mut ProgramTestContext,
+    test_data: &TestData,
+    mode: Mode,
+    token_program_id: &Pubkey,
+) {
     if mode == Mode::Burning {
         // we set the mint authority to the ntt contract in burn/mint mode
-        spl_token::instruction::set_authority(
-            &spl_token::ID,
+        spl_token_2022::instruction::set_authority(
+            token_program_id,
             &test_data.mint,
             Some(&test_data.ntt.token_authority()),
-            AuthorityType::MintTokens,
+            spl_token_2022::instruction::AuthorityType::MintTokens,
             &test_data.mint_authority.pubkey(),
             &[],
         )
@@ -142,7 +173,7 @@ pub async fn setup_ntt(ctx: &mut ProgramTestContext, test_data: &TestData, mode:
         .unwrap();
     }
 
-    initialize(
+    initialize_with_token_program_id(
         &test_data.ntt,
         Initialize {
             payer: ctx.payer.pubkey(),
@@ -155,6 +186,7 @@ pub async fn setup_ntt(ctx: &mut ProgramTestContext, test_data: &TestData, mode:
             limit: OUTBOUND_LIMIT,
             mode,
         },
+        token_program_id,
     )
     .submit_with_signers(&[&test_data.program_owner], ctx)
     .await
@@ -265,6 +297,71 @@ pub async fn setup_accounts(ctx: &mut ProgramTestContext, program_owner: Keypair
     }
 }
 
+pub async fn setup_accounts_with_transfer_fee(
+    ctx: &mut ProgramTestContext,
+    program_owner: Keypair,
+) -> TestData {
+    // create mint
+    let mint = Keypair::new();
+    let mint_authority = Keypair::new();
+
+    let user = Keypair::new();
+    let payer = ctx.payer.pubkey();
+
+    create_mint_with_transfer_fee(ctx, &mint, &mint_authority.pubkey(), 9, 500, 5000)
+        .await
+        .submit(ctx)
+        .await
+        .unwrap();
+
+    // create associated token account for user
+    let user_token_account = get_associated_token_address_with_program_id(
+        &user.pubkey(),
+        &mint.pubkey(),
+        &spl_token_2022::id(),
+    );
+
+    spl_associated_token_account::instruction::create_associated_token_account(
+        &payer,
+        &user.pubkey(),
+        &mint.pubkey(),
+        &spl_token_2022::id(),
+    )
+    .submit(ctx)
+    .await
+    .unwrap();
+
+    spl_token_2022::instruction::mint_to(
+        &spl_token_2022::id(),
+        &mint.pubkey(),
+        &user_token_account,
+        &mint_authority.pubkey(),
+        &[],
+        MINT_AMOUNT,
+    )
+    .unwrap()
+    .submit_with_signers(&[&mint_authority], ctx)
+    .await
+    .unwrap();
+
+    TestData {
+        ntt: NTT {
+            program: example_native_token_transfers::ID,
+            wormhole: Wormhole {
+                program: wormhole_anchor_sdk::wormhole::program::ID,
+            },
+        },
+        governance: Governance {
+            program: wormhole_governance::ID,
+        },
+        program_owner,
+        mint_authority,
+        mint: mint.pubkey(),
+        user,
+        user_token_account,
+    }
+}
+
 pub async fn create_mint(
     ctx: &mut ProgramTestContext,
     mint: &Keypair,
@@ -300,6 +397,57 @@ pub async fn create_mint(
     )
 }
 
+pub async fn create_mint_with_transfer_fee(
+    ctx: &mut ProgramTestContext,
+    mint: &Keypair,
+    mint_authority: &Pubkey,
+    decimals: u8,
+    transfer_fee_basis_points: u16,
+    maximum_fee: u64,
+) -> Transaction {
+    let rent = ctx.banks_client.get_rent().await.unwrap();
+    let extension_types = vec![spl_token_2022::extension::ExtensionType::TransferFeeConfig];
+    let space = spl_token_2022::extension::ExtensionType::try_calculate_account_len::<
+        spl_token_2022::state::Mint,
+    >(&extension_types)
+    .unwrap();
+    let mint_rent = rent.minimum_balance(space);
+
+    let blockhash = ctx.banks_client.get_latest_blockhash().await.unwrap();
+
+    Transaction::new_signed_with_payer(
+        &[
+            system_instruction::create_account(
+                &ctx.payer.pubkey(),
+                &mint.pubkey(),
+                mint_rent,
+                space as u64,
+                &spl_token_2022::id(),
+            ),
+            spl_token_2022::extension::transfer_fee::instruction::initialize_transfer_fee_config(
+                &spl_token_2022::id(),
+                &mint.pubkey(),
+                None,
+                None,
+                transfer_fee_basis_points,
+                maximum_fee,
+            )
+            .unwrap(),
+            spl_token_2022::instruction::initialize_mint2(
+                &spl_token_2022::id(),
+                &mint.pubkey(),
+                mint_authority,
+                None,
+                decimals,
+            )
+            .unwrap(),
+        ],
+        Some(&ctx.payer.pubkey()),
+        &[&ctx.payer, &mint],
+        blockhash,
+    )
+}
+
 // TODO: upstream this to solana-program-test
 
 /// Add a SBF program to the test environment. (copied from solana_program_test
diff --git a/solana/programs/example-native-token-transfers/tests/sdk/accounts.rs b/solana/programs/example-native-token-transfers/tests/sdk/accounts.rs
@@ -182,10 +182,18 @@ impl NTT {
     }
 
     pub fn custody(&self, mint: &Pubkey) -> Pubkey {
+        self.custody_with_token_program_id(mint, &spl_token::ID)
+    }
+
+    pub fn custody_with_token_program_id(
+        &self,
+        mint: &Pubkey,
+        token_program_id: &Pubkey,
+    ) -> Pubkey {
         anchor_spl::associated_token::get_associated_token_address_with_program_id(
             &self.token_authority(),
             mint,
-            &spl_token::ID,
+            token_program_id,
         )
     }
 
diff --git a/solana/programs/example-native-token-transfers/tests/sdk/instructions/initialize.rs b/solana/programs/example-native-token-transfers/tests/sdk/instructions/initialize.rs
@@ -13,6 +13,15 @@ pub struct Initialize {
 }
 
 pub fn initialize(ntt: &NTT, accounts: Initialize, args: InitializeArgs) -> Instruction {
+    initialize_with_token_program_id(ntt, accounts, args, &Token::id())
+}
+
+pub fn initialize_with_token_program_id(
+    ntt: &NTT,
+    accounts: Initialize,
+    args: InitializeArgs,
+    token_program_id: &Pubkey,
+) -> Instruction {
     let data = example_native_token_transfers::instruction::Initialize { args };
 
     let bpf_loader_upgradeable_program = BpfLoaderUpgradeable::id();
@@ -24,8 +33,8 @@ pub fn initialize(ntt: &NTT, accounts: Initialize, args: InitializeArgs) -> Inst
         mint: accounts.mint,
         rate_limit: ntt.outbox_rate_limit(),
         token_authority: ntt.token_authority(),
-        custody: ntt.custody(&accounts.mint),
-        token_program: Token::id(),
+        custody: ntt.custody_with_token_program_id(&accounts.mint, token_program_id),
+        token_program: *token_program_id,
         associated_token_program: AssociatedToken::id(),
         bpf_loader_upgradeable_program,
         system_program: System::id(),
diff --git a/solana/programs/example-native-token-transfers/tests/sdk/instructions/transfer.rs b/solana/programs/example-native-token-transfers/tests/sdk/instructions/transfer.rs
@@ -17,19 +17,38 @@ pub struct Transfer {
 }
 
 pub fn transfer(ntt: &NTT, transfer: Transfer, args: TransferArgs, mode: Mode) -> Instruction {
+    transfer_with_token_program_id(ntt, transfer, args, mode, &Token::id())
+}
+
+pub fn transfer_with_token_program_id(
+    ntt: &NTT,
+    transfer: Transfer,
+    args: TransferArgs,
+    mode: Mode,
+    token_program_id: &Pubkey,
+) -> Instruction {
     match mode {
-        Mode::Burning => transfer_burn(ntt, transfer, args),
-        Mode::Locking => transfer_lock(ntt, transfer, args),
+        Mode::Burning => transfer_burn_with_token_program_id(ntt, transfer, args, token_program_id),
+        Mode::Locking => transfer_lock_with_token_program_id(ntt, transfer, args, token_program_id),
     }
 }
 
 pub fn transfer_burn(ntt: &NTT, transfer: Transfer, args: TransferArgs) -> Instruction {
+    transfer_burn_with_token_program_id(ntt, transfer, args, &Token::id())
+}
+
+pub fn transfer_burn_with_token_program_id(
+    ntt: &NTT,
+    transfer: Transfer,
+    args: TransferArgs,
+    token_program_id: &Pubkey,
+) -> Instruction {
     let chain_id = args.recipient_chain.id;
     let session_authority = ntt.session_authority(&transfer.from_authority, &args);
     let data = example_native_token_transfers::instruction::TransferBurn { args };
 
     let accounts = example_native_token_transfers::accounts::TransferBurn {
-        common: common(ntt, &transfer),
+        common: common_with_token_program_id(ntt, &transfer, token_program_id),
         inbox_rate_limit: ntt.inbox_rate_limit(chain_id),
         peer: transfer.peer,
         session_authority,
@@ -44,12 +63,21 @@ pub fn transfer_burn(ntt: &NTT, transfer: Transfer, args: TransferArgs) -> Instr
 }
 
 pub fn transfer_lock(ntt: &NTT, transfer: Transfer, args: TransferArgs) -> Instruction {
+    transfer_lock_with_token_program_id(ntt, transfer, args, &Token::id())
+}
+
+pub fn transfer_lock_with_token_program_id(
+    ntt: &NTT,
+    transfer: Transfer,
+    args: TransferArgs,
+    token_program_id: &Pubkey,
+) -> Instruction {
     let chain_id = args.recipient_chain.id;
     let session_authority = ntt.session_authority(&transfer.from_authority, &args);
     let data = example_native_token_transfers::instruction::TransferLock { args };
 
     let accounts = example_native_token_transfers::accounts::TransferLock {
-        common: common(ntt, &transfer),
+        common: common_with_token_program_id(ntt, &transfer, token_program_id),
         inbox_rate_limit: ntt.inbox_rate_limit(chain_id),
         peer: transfer.peer,
         session_authority,
@@ -66,9 +94,19 @@ pub fn approve_token_authority(
     user_token_account: &Pubkey,
     user: &Pubkey,
     args: &TransferArgs,
+) -> Instruction {
+    approve_token_authority_with_token_program_id(ntt, user_token_account, user, args, &Token::id())
+}
+
+pub fn approve_token_authority_with_token_program_id(
+    ntt: &NTT,
+    user_token_account: &Pubkey,
+    user: &Pubkey,
+    args: &TransferArgs,
+    token_program_id: &Pubkey,
 ) -> Instruction {
     spl_token_2022::instruction::approve(
-        &spl_token::id(), // TODO: look into how token account was originally created
+        token_program_id,
         user_token_account,
         &ntt.session_authority(user, args),
         user,
@@ -78,18 +116,22 @@ pub fn approve_token_authority(
     .unwrap()
 }
 
-fn common(ntt: &NTT, transfer: &Transfer) -> example_native_token_transfers::accounts::Transfer {
+fn common_with_token_program_id(
+    ntt: &NTT,
+    transfer: &Transfer,
+    token_program_id: &Pubkey,
+) -> example_native_token_transfers::accounts::Transfer {
     example_native_token_transfers::accounts::Transfer {
         payer: transfer.payer,
         config: NotPausedConfig {
             config: ntt.config(),
         },
         mint: transfer.mint,
         from: transfer.from,
-        token_program: Token::id(),
+        token_program: *token_program_id,
         outbox_item: transfer.outbox_item,
         outbox_rate_limit: ntt.outbox_rate_limit(),
         system_program: System::id(),
-        custody: ntt.custody(&transfer.mint),
+        custody: ntt.custody_with_token_program_id(&transfer.mint, token_program_id),
     }
 }
diff --git a/solana/programs/example-native-token-transfers/tests/transfer.rs b/solana/programs/example-native-token-transfers/tests/transfer.rs

Original file line number	Diff line number	Diff line change
`@@ -182,10 +182,18 @@ impl NTT {`
`182`	`182`	`}`
`183`	`183`
`184`	`184`	`pub fn custody(&self, mint: &Pubkey) -> Pubkey {`
	`185`	`+ self.custody_with_token_program_id(mint, &spl_token::ID)`
	`186`	`+ }`
	`187`	`+`
	`188`	`+ pub fn custody_with_token_program_id(`
	`189`	`+ &self,`
	`190`	`+ mint: &Pubkey,`
	`191`	`+ token_program_id: &Pubkey,`
	`192`	`+ ) -> Pubkey {`
`185`	`193`	`anchor_spl::associated_token::get_associated_token_address_with_program_id(`
`186`	`194`	`&self.token_authority(),`
`187`	`195`	`mint,`
`188`		`- &spl_token::ID,`
	`196`	`+ token_program_id,`
`189`	`197`	`)`
`190`	`198`	`}`
`191`	`199`