Add Pool Test for Rust-sdk (#564)

* 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

* Initial pool test commit

* Fix bug in the pool.rs (Not getting uninitialized pool)
Update pool test with token pair
RPC update for GPA

* fmt update

* Update fmt

* Update pool test

* fmt, cleanup comments

* Restore RPC, add ignore to GPA using tests

rust-sdk/whirlpool/src/pool.rs

@@ -356,9 +356,9 @@ pub async fn fetch_whirlpools_by_token_pair(
     let whirlpool_infos = rpc.get_multiple_accounts(&whirlpool_addresses).await?;
 
     let mut whirlpools: Vec<PoolInfo> = Vec::new();
-    for i in 0..whirlpool_infos.len() {
+    for i in 0..whirlpool_addresses.len() {
         let pool_address = whirlpool_addresses[i];
-        let pool_info = whirlpool_infos[i].as_ref();
+        let pool_info = whirlpool_infos.get(i).and_then(|x| x.as_ref());
         let fee_tier = &fee_tiers[i];
 
         if let Some(pool_info) = pool_info {
@@ -380,3 +380,170 @@ pub async fn fetch_whirlpools_by_token_pair(
 
     Ok(whirlpools)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::tests::{
+        setup_ata_with_amount, setup_mint_with_decimals, setup_whirlpool, RpcContext,
+    };
+    use serial_test::serial;
+    use solana_program_test::tokio;
+    use solana_sdk::signer::Signer;
+
+    struct TestContext {
+        ctx: RpcContext,
+        mint_a: Pubkey,
+        mint_b: Pubkey,
+        splash_pool: Pubkey,
+        concentrated_pool: Pubkey,
+    }
+
+    impl TestContext {
+        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?;
+
+            setup_ata_with_amount(&ctx, mint_a, 500_000_000_000).await?;
+            setup_ata_with_amount(&ctx, mint_b, 500_000_000_000).await?;
+
+            // Setup all pools
+            let concentrated_pool = setup_whirlpool(&ctx, mint_a, mint_b, 64).await?;
+            let splash_pool =
+                setup_whirlpool(&ctx, mint_a, mint_b, SPLASH_POOL_TICK_SPACING).await?;
+
+            Ok(Self {
+                ctx,
+                mint_a,
+                mint_b,
+                splash_pool,
+                concentrated_pool,
+            })
+        }
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn test_fetch_splash_pool() {
+        let test_ctx = TestContext::new().await.unwrap();
+
+        if let PoolInfo::Initialized(pool) =
+            fetch_splash_pool(&test_ctx.ctx.rpc, test_ctx.mint_a, test_ctx.mint_b)
+                .await
+                .unwrap()
+        {
+            assert_eq!(pool.data.liquidity, 0);
+            assert_eq!(pool.data.tick_spacing, SPLASH_POOL_TICK_SPACING);
+            assert_eq!(pool.address, test_ctx.splash_pool);
+            assert_eq!(pool.data.token_mint_a, test_ctx.mint_a);
+            assert_eq!(pool.data.token_mint_b, test_ctx.mint_b);
+            assert_eq!(pool.data.fee_rate, 1000);
+            assert_eq!(pool.data.protocol_fee_rate, 0);
+        } else {
+            panic!("Expected initialized pool");
+        }
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn test_fetch_concentrated_liquidity_pool() {
+        let test_ctx = TestContext::new().await.unwrap();
+
+        if let PoolInfo::Initialized(pool) = fetch_concentrated_liquidity_pool(
+            &test_ctx.ctx.rpc,
+            test_ctx.mint_a,
+            test_ctx.mint_b,
+            64,
+        )
+        .await
+        .unwrap()
+        {
+            assert_eq!(pool.data.liquidity, 0);
+            assert_eq!(pool.data.tick_spacing, 64);
+            assert_eq!(pool.address, test_ctx.concentrated_pool);
+            assert_eq!(pool.data.token_mint_a, test_ctx.mint_a);
+            assert_eq!(pool.data.token_mint_b, test_ctx.mint_b);
+            assert_eq!(pool.data.fee_rate, 300);
+            assert_eq!(pool.data.protocol_fee_rate, 0);
+        } else {
+            panic!("Expected initialized pool");
+        }
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn test_fetch_non_existent_pool() {
+        let test_ctx = TestContext::new().await.unwrap();
+
+        if let PoolInfo::Uninitialized(pool) = fetch_concentrated_liquidity_pool(
+            &test_ctx.ctx.rpc,
+            test_ctx.mint_a,
+            test_ctx.mint_b,
+            128,
+        )
+        .await
+        .unwrap()
+        {
+            assert_eq!(pool.tick_spacing, 128);
+            assert_eq!(pool.token_mint_a, test_ctx.mint_a);
+            assert_eq!(pool.token_mint_b, test_ctx.mint_b);
+            assert_eq!(pool.fee_rate, 1000);
+            assert_eq!(pool.protocol_fee_rate, 0);
+        } else {
+            panic!("Expected uninitialized pool");
+        }
+    }
+
+    #[tokio::test]
+    #[serial]
+    #[ignore = "Skipped until solana-bankrun supports getProgramAccounts"]
+    async fn test_fetch_all_pools_for_pair() {
+        let test_ctx = TestContext::new().await.unwrap();
+
+        let pools =
+            fetch_whirlpools_by_token_pair(&test_ctx.ctx.rpc, test_ctx.mint_a, test_ctx.mint_b)
+                .await
+                .unwrap();
+
+        assert_eq!(pools.len(), 3); // 2 initialized + 1 uninitialized (128 tick spacing)
+
+        // Verify concentrated liquidity pool
+        let concentrated = pools
+            .iter()
+            .find(|p| match p {
+                PoolInfo::Initialized(p) => p.data.tick_spacing == 64,
+                _ => false,
+            })
+            .unwrap();
+        if let PoolInfo::Initialized(pool) = concentrated {
+            assert_eq!(pool.address, test_ctx.concentrated_pool);
+            assert_eq!(pool.data.fee_rate, 300);
+        }
+
+        // Verify splash pool
+        let splash = pools
+            .iter()
+            .find(|p| match p {
+                PoolInfo::Initialized(p) => p.data.tick_spacing == SPLASH_POOL_TICK_SPACING,
+                _ => false,
+            })
+            .unwrap();
+        if let PoolInfo::Initialized(pool) = splash {
+            assert_eq!(pool.address, test_ctx.splash_pool);
+            assert_eq!(pool.data.fee_rate, 1000);
+        }
+
+        // Verify uninitialized pool
+        let uninitialized = pools
+            .iter()
+            .find(|p| match p {
+                PoolInfo::Uninitialized(p) => p.tick_spacing == 128,
+                _ => false,
+            })
+            .unwrap();
+        if let PoolInfo::Uninitialized(pool) = uninitialized {
+            assert_eq!(pool.fee_rate, 1000);
+        }
+    }
+}

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

@@ -5,6 +5,7 @@ mod token;
 mod token_extensions;
 
 pub use anchor::*;
+pub use program::*;
 pub use rpc::*;
 pub use token::*;
 pub use token_extensions::*;

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

@@ -1,8 +1,63 @@
-use solana_sdk::pubkey::Pubkey;
+use solana_sdk::{pubkey::Pubkey, signer::Signer, system_program};
 use std::error::Error;
 
-pub async fn setup_whirlpool() -> Result<Pubkey, Box<dyn Error>> {
-    todo!()
+use orca_whirlpools_client::{
+    get_fee_tier_address, get_token_badge_address, get_whirlpool_address, InitializePoolV2,
+    InitializePoolV2InstructionArgs,
+};
+use orca_whirlpools_core::{price_to_sqrt_price, tick_index_to_sqrt_price};
+use solana_program::sysvar::rent::ID as RENT_PROGRAM_ID;
+
+use crate::WHIRLPOOLS_CONFIG_ADDRESS;
+
+use super::rpc::RpcContext;
+
+pub async fn setup_whirlpool(
+    ctx: &RpcContext,
+    token_a: Pubkey,
+    token_b: Pubkey,
+    tick_spacing: u16,
+) -> Result<Pubkey, Box<dyn Error>> {
+    let config = *WHIRLPOOLS_CONFIG_ADDRESS.try_lock()?;
+    let fee_tier = get_fee_tier_address(&config, tick_spacing)?.0;
+    let whirlpool = get_whirlpool_address(&config, &token_a, &token_b, tick_spacing)?.0;
+    let token_badge_a = get_token_badge_address(&config, &token_a)?.0;
+    let token_badge_b = get_token_badge_address(&config, &token_b)?.0;
+
+    let vault_a = ctx.get_next_keypair();
+    let vault_b = ctx.get_next_keypair();
+
+    let mint_a_info = ctx.rpc.get_account(&token_a).await?;
+    let mint_b_info = ctx.rpc.get_account(&token_b).await?;
+
+    // Default initial price of 1.0
+    let sqrt_price = tick_index_to_sqrt_price(0);
+
+    let instructions = vec![InitializePoolV2 {
+        whirlpool,
+        fee_tier,
+        token_mint_a: token_a,
+        token_mint_b: token_b,
+        whirlpools_config: config,
+        funder: ctx.signer.pubkey(),
+        token_vault_a: vault_a.pubkey(),
+        token_vault_b: vault_b.pubkey(),
+        token_badge_a,
+        token_badge_b,
+        token_program_a: mint_a_info.owner,
+        token_program_b: mint_b_info.owner,
+        system_program: system_program::id(),
+        rent: RENT_PROGRAM_ID,
+    }
+    .instruction(InitializePoolV2InstructionArgs {
+        tick_spacing,
+        initial_sqrt_price: sqrt_price,
+    })];
+
+    ctx.send_transaction_with_signers(instructions, vec![&vault_a, &vault_b])
+        .await?;
+
+    Ok(whirlpool)
 }
 
 pub async fn setup_position(_whirlpool: Pubkey) -> Result<Pubkey, Box<dyn Error>> {