Add Swap test for Rust SDK

rust-sdk/whirlpool/src/harvest.rs

@@ -367,6 +367,7 @@ mod tests {
 
     async fn get_token_balance(rpc: &RpcClient, address: Pubkey) -> Result<u64, Box<dyn Error>> {
         let account_data = rpc.get_account(&address).await?;
+
         if account_data.owner == TOKEN_2022_PROGRAM_ID {
             let parsed = StateWithExtensionsOwned::<TokenAccount2022>::unpack(account_data.data)?;
             Ok(parsed.base.amount)
@@ -381,6 +382,7 @@ mod tests {
         harvest_ix: &HarvestPositionInstruction,
         ata_a: Pubkey,
         ata_b: Pubkey,
+
         position_mint: Pubkey,
     ) -> Result<(), Box<dyn Error>> {
         let before_a = get_token_balance(&ctx.rpc, ata_a).await?;

rust-sdk/whirlpool/src/swap.rs

@@ -323,3 +323,267 @@ pub async fn swap_instructions(
         additional_signers: token_accounts.additional_signers,
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use serial_test::serial;
+    use solana_client::nonblocking::rpc_client::RpcClient;
+    use solana_sdk::{program_pack::Pack, pubkey::Pubkey, signer::Signer};
+    use spl_token::state::Account as TokenAccount;
+    use spl_token_2022::state::Account as TokenAccount2022;
+    use spl_token_2022::{extension::StateWithExtensionsOwned, ID as TOKEN_2022_PROGRAM_ID};
+    use std::error::Error;
+
+    use crate::{
+        create_concentrated_liquidity_pool_instructions, create_splash_pool_instructions,
+        increase_liquidity_instructions, open_full_range_position_instructions, swap_instructions,
+        tests::{setup_ata_with_amount, setup_mint_with_decimals, RpcContext},
+        IncreaseLiquidityParam, SwapQuote, SwapType,
+    };
+
+    struct SwapTestContext {
+        pub ctx: RpcContext,
+
+        pub mint_a: Pubkey,
+        pub mint_b: Pubkey,
+        pub ata_a: Pubkey,
+        pub ata_b: Pubkey,
+    }
+
+    impl SwapTestContext {
+        pub async fn new() -> Result<Self, Box<dyn Error>> {
+            let ctx = RpcContext::new().await;
+
+            let mint_a = setup_mint_with_decimals(&ctx, 9).await?;
+            let mint_b = setup_mint_with_decimals(&ctx, 9).await?;
+
+            let ata_a = setup_ata_with_amount(&ctx, mint_a, 1_000_000_000).await?;
+            let ata_b = setup_ata_with_amount(&ctx, mint_b, 1_000_000_000).await?;
+
+            Ok(Self {
+                ctx,
+                mint_a,
+                mint_b,
+                ata_a,
+                ata_b,
+            })
+        }
+
+        async fn get_token_balance(&self, address: Pubkey) -> Result<u64, Box<dyn Error>> {
+            let account_data = self.ctx.rpc.get_account(&address).await?;
+            if account_data.owner == TOKEN_2022_PROGRAM_ID {
+                let parsed =
+                    StateWithExtensionsOwned::<TokenAccount2022>::unpack(account_data.data)?;
+                Ok(parsed.base.amount)
+            } else {
+                let parsed = TokenAccount::unpack(&account_data.data)?;
+                Ok(parsed.amount)
+            }
+        }
+
+        pub async fn init_pool(&self, is_splash: bool) -> Result<Pubkey, Box<dyn Error>> {
+            if is_splash {
+                let pool = create_splash_pool_instructions(
+                    &self.ctx.rpc,
+                    self.mint_a,
+                    self.mint_b,
+                    None,
+                    Some(self.ctx.signer.pubkey()),
+                )
+                .await?;
+                self.ctx
+                    .send_transaction_with_signers(
+                        pool.instructions,
+                        pool.additional_signers.iter().collect(),
+                    )
+                    .await?;
+                Ok(pool.pool_address)
+            } else {
+                let cl_pool = create_concentrated_liquidity_pool_instructions(
+                    &self.ctx.rpc,
+                    self.mint_a,
+                    self.mint_b,
+                    128,
+                    None,
+                    Some(self.ctx.signer.pubkey()),
+                )
+                .await?;
+                self.ctx
+                    .send_transaction_with_signers(
+                        cl_pool.instructions,
+                        cl_pool.additional_signers.iter().collect(),
+                    )
+                    .await?;
+                Ok(cl_pool.pool_address)
+            }
+        }
+
+        pub async fn open_position_with_liquidity(
+            &self,
+            pool_pubkey: Pubkey,
+        ) -> Result<Pubkey, Box<dyn Error>> {
+            let position = open_full_range_position_instructions(
+                &self.ctx.rpc,
+                pool_pubkey,
+                IncreaseLiquidityParam::Liquidity(50_000_000),
+                None,
+                Some(self.ctx.signer.pubkey()),
+            )
+            .await?;
+            self.ctx
+                .send_transaction_with_signers(
+                    position.instructions,
+                    position.additional_signers.iter().collect(),
+                )
+                .await?;
+
+            Ok(position.position_mint)
+        }
+
+        pub async fn do_swap(
+            &self,
+            pool_pubkey: Pubkey,
+            a_to_b: bool,
+            swap_type: SwapType,
+            amount: u64,
+        ) -> Result<(), Box<dyn Error>> {
+            let specified_mint = if a_to_b { self.mint_a } else { self.mint_b };
+
+            let before_a = self.get_token_balance(self.ata_a).await?;
+            let before_b = self.get_token_balance(self.ata_b).await?;
+
+            let swap_ix = swap_instructions(
+                &self.ctx.rpc,
+                pool_pubkey,
+                amount,
+                specified_mint,
+                swap_type.clone(),
+                Some(100), // 1% slippage
+                Some(self.ctx.signer.pubkey()),
+            )
+            .await?;
+
+            self.ctx
+                .send_transaction_with_signers(
+                    swap_ix.instructions,
+                    swap_ix.additional_signers.iter().collect(),
+                )
+                .await?;
+
+            let after_a = self.get_token_balance(self.ata_a).await?;
+            let after_b = self.get_token_balance(self.ata_b).await?;
+
+            let used_a = before_a.saturating_sub(after_a);
+            let used_b = before_b.saturating_sub(after_b);
+            let gained_a = after_a.saturating_sub(before_a);
+            let gained_b = after_b.saturating_sub(before_b);
+
+            match swap_ix.quote {
+                SwapQuote::ExactIn(q) => {
+                    if a_to_b {
+                        // used A, gained B
+                        assert_eq!(used_a, q.token_in, "Used A mismatch");
+                        assert_eq!(gained_b, q.token_est_out, "Gained B mismatch");
+                    } else {
+                        // used B, gained A
+                        assert_eq!(used_b, q.token_in, "Used B mismatch");
+                        assert_eq!(gained_a, q.token_est_out, "Gained A mismatch");
+                    }
+                }
+                SwapQuote::ExactOut(q) => {
+                    if a_to_b {
+                        // gained B, used A
+                        assert_eq!(gained_b, q.token_out, "Gained B mismatch");
+                        assert_eq!(used_a, q.token_est_in, "Used A mismatch");
+                    } else {
+                        // gained A, used B
+                        assert_eq!(gained_a, q.token_out, "Gained A mismatch");
+                        assert_eq!(used_b, q.token_est_in, "Used B mismatch");
+                    }
+                }
+            }
+            println!(
+                "swap result => a_to_b={}, used_a={}, used_b={}, gained_a={}, gained_b={}",
+                a_to_b, used_a, used_b, gained_a, gained_b
+            );
+
+            Ok(())
+        }
+    }
+
+    #[tokio::test]
+    async fn test_swap_for_multiple_pools() -> Result<(), Box<dyn Error>> {
+        let stx = SwapTestContext::new().await?;
+
+        let ctx = &stx.ctx;
+
+        let mint_a = setup_mint_with_decimals(&ctx, 9).await?;
+        let mint_b = setup_mint_with_decimals(&ctx, 9).await?;
+        let ata_a = setup_ata_with_amount(&ctx, mint_a, 500_000_000).await?;
+        let ata_b = setup_ata_with_amount(&ctx, mint_b, 500_000_000).await?;
+
+        let pool = create_concentrated_liquidity_pool_instructions(
+            &ctx.rpc,
+            mint_a,
+            mint_b,
+            128,
+            None,
+            Some(ctx.signer.pubkey()),
+        )
+        .await?;
+        ctx.send_transaction_with_signers(
+            pool.instructions,
+            pool.additional_signers.iter().collect(),
+        )
+        .await?;
+
+        let pool_pubkey = pool.pool_address;
+
+        let position = open_full_range_position_instructions(
+            &ctx.rpc,
+            pool_pubkey,
+            IncreaseLiquidityParam::Liquidity(50_000_000),
+            None,
+            Some(ctx.signer.pubkey()),
+        )
+        .await?;
+        ctx.send_transaction_with_signers(
+            position.instructions,
+            position.additional_signers.iter().collect(),
+        )
+        .await?;
+
+        let swap_ix = swap_instructions(
+            &ctx.rpc,
+            pool_pubkey,
+            10_000,
+            mint_a,
+            SwapType::ExactIn,
+            Some(100),
+            Some(ctx.signer.pubkey()),
+        )
+        .await?;
+        let before_a = stx.get_token_balance(ata_a).await?;
+        let before_b = stx.get_token_balance(ata_b).await?;
+
+        ctx.send_transaction_with_signers(
+            swap_ix.instructions,
+            swap_ix.additional_signers.iter().collect(),
+        )
+        .await?;
+
+        let after_a = stx.get_token_balance(ata_a).await?;
+        let after_b = stx.get_token_balance(ata_b).await?;
+        let used_a = before_a.saturating_sub(after_a);
+        let gained_b = after_b.saturating_sub(before_b);
+
+        if let SwapQuote::ExactIn(q) = swap_ix.quote {
+            assert_eq!(used_a, q.token_in, "Used A mismatch");
+            assert_eq!(gained_b, q.token_est_out, "Gained B mismatch");
+        } else {
+            panic!("Expected ExactIn quote");
+        }
+
+        Ok(())
+    }
+}