Add Token Testing for Rust Sdk (#563)

* Initial commit for token test for rust sdk. Reorganize test files to utils

* Cleanup file path

* fmt fix

* Revert utils folder, Add token tests, Update extension test to check variables in the token 2022 not just length

* fmt, update more tests

* Update test to do more checks not only len, test to check token.rs

* fmt

* Update and add more tests

* Delete useless asserts

* Remove len based asserts and do more sophisticated test

* Update native sol-wsol test and update loop match

rust-sdk/whirlpool/Cargo.toml

@@ -36,3 +36,5 @@ solana-version = { version = "^1.18" }
 async-trait = { version = "^0.1" }
 base64 = { version = "^0.20" }
 toml = { version = "^0.7" }
+tokio = { version = "1.0", features = ["sync"] }
+spl-pod = "0.1.0"

rust-sdk/whirlpool/src/tests/mod.rs

@@ -7,3 +7,4 @@ mod token_extensions;
 pub use anchor::*;
 pub use rpc::*;
 pub use token::*;
+pub use token_extensions::*;

rust-sdk/whirlpool/src/tests/token_extensions.rs

@@ -1,19 +1,134 @@
+use solana_sdk::{
+    pubkey::Pubkey,
+    signer::{keypair::Keypair, Signer},
+    system_instruction,
+};
+use spl_associated_token_account::{
+    get_associated_token_address_with_program_id,
+    instruction::create_associated_token_account_idempotent,
+};
+use spl_token_2022::{
+    extension::{
+        transfer_fee::instruction::{initialize_transfer_fee_config, set_transfer_fee},
+        ExtensionType,
+    },
+    instruction::{initialize_mint2, mint_to},
+    state::Mint,
+    ID as TOKEN_2022_PROGRAM_ID,
+};
 use std::error::Error;
 
-use solana_sdk::pubkey::Pubkey;
+use super::rpc::RpcContext;
 
-pub async fn setup_ata_te(_mint: Pubkey) -> Result<Pubkey, Box<dyn Error>> {
-    todo!()
+#[derive(Default)]
+pub struct SetupAtaConfig {
+    pub amount: Option<u64>,
 }
 
-pub async fn setup_mint_te() -> Result<Pubkey, Box<dyn Error>> {
-    todo!()
+pub async fn setup_mint_te(
+    ctx: &RpcContext,
+    extensions: &[ExtensionType],
+) -> Result<Pubkey, Box<dyn Error>> {
+    let mint = Keypair::new();
+    let mut instructions = vec![];
+
+    // 1. Create account instruction
+    let space = ExtensionType::try_calculate_account_len::<Mint>(extensions)?;
+    let rent = ctx
+        .rpc
+        .get_minimum_balance_for_rent_exemption(space)
+        .await?;
+
+    instructions.push(system_instruction::create_account(
+        &ctx.signer.pubkey(),
+        &mint.pubkey(),
+        rent,
+        space as u64,
+        &TOKEN_2022_PROGRAM_ID,
+    ));
+
+    // 2. Initialize extensions first
+    for extension in extensions {
+        match extension {
+            ExtensionType::TransferFeeConfig => {
+                instructions.push(initialize_transfer_fee_config(
+                    &TOKEN_2022_PROGRAM_ID,
+                    &mint.pubkey(),
+                    Some(&ctx.signer.pubkey()),
+                    Some(&ctx.signer.pubkey()),
+                    100,           // 1% (matching program)
+                    1_000_000_000, // 1 token (matching program)
+                )?);
+            }
+            _ => {} // Handle other extension types here if needed
+        }
+    }
+
+    // 3. Initialize mint
+    instructions.push(initialize_mint2(
+        &TOKEN_2022_PROGRAM_ID,
+        &mint.pubkey(),
+        &ctx.signer.pubkey(),
+        None, // freeze_authority
+        6,    // decimals
+    )?);
+
+    // 4. Set extension configurations
+    for extension in extensions {
+        match extension {
+            ExtensionType::TransferFeeConfig => {
+                instructions.push(set_transfer_fee(
+                    &TOKEN_2022_PROGRAM_ID,
+                    &mint.pubkey(),
+                    &ctx.signer.pubkey(),
+                    &[],
+                    150,           // 1.5% (matching program)
+                    1_000_000_000, // 1 token
+                )?);
+            }
+            _ => {} // Handle other extension types here if needed
+        }
+    }
+
+    ctx.send_transaction_with_signers(instructions, vec![&mint])
+        .await?;
+    Ok(mint.pubkey())
 }
 
-pub async fn setup_mint_te_fee() -> Result<Pubkey, Box<dyn Error>> {
-    todo!()
+pub async fn setup_mint_te_fee(ctx: &RpcContext) -> Result<Pubkey, Box<dyn Error>> {
+    setup_mint_te(ctx, &[ExtensionType::TransferFeeConfig]).await
 }
 
-pub async fn setup_mint_te_hook() -> Result<Pubkey, Box<dyn Error>> {
-    todo!()
+pub async fn setup_ata_te(
+    ctx: &RpcContext,
+    mint: Pubkey,
+    config: Option<SetupAtaConfig>,
+) -> Result<Pubkey, Box<dyn Error>> {
+    let config = config.unwrap_or_default();
+    let ata = get_associated_token_address_with_program_id(
+        &ctx.signer.pubkey(),
+        &mint,
+        &TOKEN_2022_PROGRAM_ID,
+    );
+
+    let mut instructions = vec![create_associated_token_account_idempotent(
+        &ctx.signer.pubkey(),
+        &ctx.signer.pubkey(),
+        &mint,
+        &TOKEN_2022_PROGRAM_ID,
+    )];
+
+    if let Some(amount) = config.amount {
+        instructions.push(mint_to(
+            &TOKEN_2022_PROGRAM_ID,
+            &mint,
+            &ata,
+            &ctx.signer.pubkey(),
+            &[],
+            amount,
+        )?);
+    }
+
+    ctx.send_transaction(instructions).await?;
+    Ok(ata)
 }