From ead6aed438814c434cf34689dbda1b7a6ddf6159 Mon Sep 17 00:00:00 2001 From: pauldragonfly Date: Fri, 17 Jan 2025 00:02:07 +0900 Subject: [PATCH] Add decrease liquidity test for rust sdk (#608) * Add decrease liquidity test * fmt * Update comments * Fix fmt * Cleanup repo * Cleanup * Update multiple combo of tests Update program to handle multiple tick array when decrease --- rust-sdk/whirlpool/src/decrease_liquidity.rs | 223 ++++++++++++++++++- rust-sdk/whirlpool/src/tests/program.rs | 87 +++++++- 2 files changed, 298 insertions(+), 12 deletions(-) diff --git a/rust-sdk/whirlpool/src/decrease_liquidity.rs b/rust-sdk/whirlpool/src/decrease_liquidity.rs index 15a2c99fb..59b18ba4e 100644 --- a/rust-sdk/whirlpool/src/decrease_liquidity.rs +++ b/rust-sdk/whirlpool/src/decrease_liquidity.rs @@ -1,9 +1,3 @@ -use std::{ - collections::HashSet, - error::Error, - time::{SystemTime, UNIX_EPOCH}, -}; - use orca_whirlpools_client::{ get_position_address, get_tick_array_address, Position, TickArray, Whirlpool, }; @@ -18,8 +12,14 @@ use orca_whirlpools_core::{ get_tick_index_in_array, CollectFeesQuote, CollectRewardsQuote, DecreaseLiquidityQuote, }; use solana_client::nonblocking::rpc_client::RpcClient; +use solana_sdk::program_pack::Pack; use solana_sdk::{account::Account, instruction::Instruction, pubkey::Pubkey, signature::Keypair}; use spl_associated_token_account::get_associated_token_address_with_program_id; +use std::{ + collections::HashSet, + error::Error, + time::{SystemTime, UNIX_EPOCH}, +}; use crate::{ token::{get_current_transfer_fee, prepare_token_accounts_instructions, TokenAccountStrategy}, @@ -635,3 +635,214 @@ pub async fn close_position_instructions( rewards_quote, }) } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::error::Error; + + use serial_test::serial; + use solana_program_test::tokio; + use solana_sdk::{ + program_pack::Pack, + pubkey::Pubkey, + signer::{keypair::Keypair, Signer}, + }; + + use solana_client::nonblocking::rpc_client::RpcClient; + use spl_token::state::Account as TokenAccount; + use spl_token_2022::{ + extension::StateWithExtensionsOwned, state::Account as TokenAccount2022, + ID as TOKEN_2022_PROGRAM_ID, + }; + + use crate::{ + decrease_liquidity_instructions, increase_liquidity_instructions, + tests::{ + setup_ata_te, setup_ata_with_amount, setup_mint_te, setup_mint_te_fee, + setup_mint_with_decimals, setup_position, setup_whirlpool, RpcContext, SetupAtaConfig, + }, + DecreaseLiquidityParam, IncreaseLiquidityParam, + }; + use orca_whirlpools_client::{get_position_address, Position}; + + async fn get_token_balance(rpc: &RpcClient, address: Pubkey) -> Result> { + let account_data = rpc.get_account(&address).await?; + + if account_data.owner == TOKEN_2022_PROGRAM_ID { + // Token-2022 + (possibly) extension + let state = StateWithExtensionsOwned::::unpack(account_data.data)?; + Ok(state.base.amount) + } else { + // SPL Token + let token_account = TokenAccount::unpack(&account_data.data)?; + Ok(token_account.amount) + } + } + + async fn fetch_position( + rpc: &solana_client::nonblocking::rpc_client::RpcClient, + position_pubkey: Pubkey, + ) -> Result> { + let account = rpc.get_account(&position_pubkey).await?; + Ok(Position::from_bytes(&account.data)?) + } + + async fn verify_decrease_liquidity( + ctx: &RpcContext, + decrease_ix: &crate::DecreaseLiquidityInstruction, + token_a_account: Pubkey, + token_b_account: Pubkey, + position_mint: Pubkey, + ) -> Result<(), Box> { + let before_a = get_token_balance(&ctx.rpc, token_a_account).await?; + let before_b = get_token_balance(&ctx.rpc, token_b_account).await?; + + let signers: Vec<&Keypair> = decrease_ix.additional_signers.iter().collect(); + ctx.send_transaction_with_signers(decrease_ix.instructions.clone(), signers) + .await?; + + let after_a = get_token_balance(&ctx.rpc, token_a_account).await?; + let after_b = get_token_balance(&ctx.rpc, token_b_account).await?; + let gained_a = after_a.saturating_sub(before_a); + let gained_b = after_b.saturating_sub(before_b); + + let quote = &decrease_ix.quote; + assert!( + gained_a >= quote.token_min_a && gained_a <= quote.token_est_a, + "Token A gain out of range: gained={}, expected={}..{}", + gained_a, + quote.token_min_a, + quote.token_est_a + ); + assert!( + gained_b >= quote.token_min_b && gained_b <= quote.token_est_b, + "Token B gain out of range: gained={}, expected={}..{}", + gained_b, + quote.token_min_b, + quote.token_est_b + ); + + let position_pubkey = get_position_address(&position_mint)?.0; + let position_data = fetch_position(&ctx.rpc, position_pubkey).await?; + assert_eq!( + position_data.liquidity, quote.liquidity_delta, + "Position liquidity mismatch! expected={}, got={}", + quote.liquidity_delta, position_data.liquidity + ); + + Ok(()) + } + + #[tokio::test] + #[serial] + async fn test_decrease_liquidity_multiple_combos() -> Result<(), Box> { + 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 mint_te_a = setup_mint_te(&ctx, &[]).await?; + let mint_te_b = setup_mint_te(&ctx, &[]).await?; + let mint_te_fee = setup_mint_te_fee(&ctx).await?; + + let token_balance = 1_000_000; + let user_ata_a = setup_ata_with_amount(&ctx, mint_a, token_balance).await?; + let user_ata_b = setup_ata_with_amount(&ctx, mint_b, token_balance).await?; + let user_ata_te_a = setup_ata_te( + &ctx, + mint_te_a, + Some(SetupAtaConfig { + amount: Some(token_balance), + }), + ) + .await?; + let user_ata_te_b = setup_ata_te( + &ctx, + mint_te_b, + Some(SetupAtaConfig { + amount: Some(token_balance), + }), + ) + .await?; + let user_ata_tefee = setup_ata_te( + &ctx, + mint_te_fee, + Some(SetupAtaConfig { + amount: Some(token_balance), + }), + ) + .await?; + + let mut minted: HashMap<&str, Pubkey> = HashMap::new(); + minted.insert("A", mint_a); + minted.insert("B", mint_b); + minted.insert("TEA", mint_te_a); + minted.insert("TEB", mint_te_b); + minted.insert("TEFee", mint_te_fee); + + let mut user_atas: HashMap<&str, Pubkey> = HashMap::new(); + user_atas.insert("A", user_ata_a); + user_atas.insert("B", user_ata_b); + user_atas.insert("TEA", user_ata_te_a); + user_atas.insert("TEB", user_ata_te_b); + user_atas.insert("TEFee", user_ata_tefee); + + let pool_combos = vec![ + ("A-B", ("A", "B")), + ("A-TEA", ("A", "TEA")), + ("TEA-TEB", ("TEA", "TEB")), + ("A-TEFee", ("A", "TEFee")), + ]; + + let position_ranges = vec![ + ("equally centered", (-100, 100)), + ("one sided A", (-100, -1)), + ("one sided B", (1, 100)), + ]; + + let tick_spacing = 64; + + for (pool_name, (mint_a_key, mint_b_key)) in &pool_combos { + let pubkey_a = *minted.get(mint_a_key).unwrap(); + let pubkey_b = *minted.get(mint_b_key).unwrap(); + + let pool_pubkey = setup_whirlpool(&ctx, pubkey_a, pubkey_b, tick_spacing).await?; + + for (range_name, (lower, upper)) in &position_ranges { + let position_mint = + setup_position(&ctx, pool_pubkey, Some((*lower, *upper)), None).await?; + + let increase_ix = increase_liquidity_instructions( + &ctx.rpc, + position_mint, + IncreaseLiquidityParam::Liquidity(100_000), + Some(100), // 1% slippage + Some(ctx.signer.pubkey()), + ) + .await?; + ctx.send_transaction_with_signers(increase_ix.instructions, vec![]) + .await?; + + let decrease_ix = decrease_liquidity_instructions( + &ctx.rpc, + position_mint, + DecreaseLiquidityParam::Liquidity(50_000), + Some(100), // 1% slippage + Some(ctx.signer.pubkey()), + ) + .await?; + + let ata_a = *user_atas.get(mint_a_key).unwrap(); + let ata_b = *user_atas.get(mint_b_key).unwrap(); + verify_decrease_liquidity(&ctx, &decrease_ix, ata_a, ata_b, position_mint).await?; + + println!( + "[combo={}, range={}, pos={}] Decrease 50k => OK", + pool_name, range_name, position_mint + ); + } + } + + Ok(()) + } +} diff --git a/rust-sdk/whirlpool/src/tests/program.rs b/rust-sdk/whirlpool/src/tests/program.rs index f9fe47be3..1610a1cce 100644 --- a/rust-sdk/whirlpool/src/tests/program.rs +++ b/rust-sdk/whirlpool/src/tests/program.rs @@ -4,10 +4,12 @@ use orca_whirlpools_client::{ get_whirlpool_address, InitializePoolV2, InitializePoolV2InstructionArgs, InitializePositionBundle, InitializeTickArray, InitializeTickArrayInstructionArgs, OpenBundledPosition, OpenBundledPositionInstructionArgs, OpenPosition, - OpenPositionInstructionArgs, Whirlpool, + OpenPositionInstructionArgs, OpenPositionWithTokenExtensions, + OpenPositionWithTokenExtensionsInstructionArgs, Whirlpool, }; use orca_whirlpools_core::{ get_initializable_tick_index, get_tick_array_start_tick_index, tick_index_to_sqrt_price, + TICK_ARRAY_SIZE, }; use solana_program::program_pack::Pack; use solana_program::sysvar::rent::ID as RENT_PROGRAM_ID; @@ -32,6 +34,61 @@ use super::rpc::RpcContext; use crate::tests::token_extensions::setup_mint_te; +pub async fn init_tick_arrays_for_range( + ctx: &RpcContext, + whirlpool: Pubkey, + lower_tick_index: i32, + upper_tick_index: i32, + spacing: u16, +) -> Result<(), Box> { + let (low, high) = if lower_tick_index <= upper_tick_index { + (lower_tick_index, upper_tick_index) + } else { + (upper_tick_index, lower_tick_index) + }; + + let offset = (spacing as i32) * (TICK_ARRAY_SIZE as i32); + + let start_low = get_tick_array_start_tick_index(low, spacing); + let start_high = get_tick_array_start_tick_index(high, spacing); + + let (begin, end) = if start_low <= start_high { + (start_low, start_high) + } else { + (start_high, start_low) + }; + + let mut instructions = vec![]; + + let mut current = begin; + while current <= end { + let (tick_array_addr, _) = get_tick_array_address(&whirlpool, current)?; + + let account_result = ctx.rpc.get_account(&tick_array_addr).await; + if account_result.is_err() { + instructions.push( + InitializeTickArray { + whirlpool, + funder: ctx.signer.pubkey(), + tick_array: tick_array_addr, + system_program: system_program::id(), + } + .instruction(InitializeTickArrayInstructionArgs { + start_tick_index: current, + }), + ); + } + + current += offset; + } + + if !instructions.is_empty() { + ctx.send_transaction(instructions).await?; + } + + Ok(()) +} + pub async fn setup_whirlpool( ctx: &RpcContext, token_a: Pubkey, @@ -106,6 +163,15 @@ pub async fn setup_position( let (lower_tick_array_addr, _) = get_tick_array_address(&whirlpool, lower_tick_array_start)?; let (upper_tick_array_addr, _) = get_tick_array_address(&whirlpool, upper_tick_array_start)?; + init_tick_arrays_for_range( + ctx, + whirlpool, + tick_lower, + tick_upper, + whirlpool_account.tick_spacing, + ) + .await?; + let mut instructions = vec![]; let lower_tick_array_account = ctx.rpc.get_account(&lower_tick_array_addr).await; @@ -198,6 +264,7 @@ pub async fn setup_te_position( tick_range: Option<(i32, i32)>, owner: Option, ) -> Result> { + let metadata_update_auth = Pubkey::try_from("3axbTs2z5GBy6usVbNVoqEgZMng3vZvMnAoX29BFfwhr")?; let owner = owner.unwrap_or_else(|| ctx.signer.pubkey()); let whirlpool_data = ctx.rpc.get_account(&whirlpool).await?; let whirlpool_account = Whirlpool::from_bytes(&whirlpool_data.data)?; @@ -215,6 +282,14 @@ pub async fn setup_te_position( get_tick_array_start_tick_index(lower_tick_index, whirlpool_account.tick_spacing), get_tick_array_start_tick_index(upper_tick_index, whirlpool_account.tick_spacing), ]; + init_tick_arrays_for_range( + ctx, + whirlpool, + tick_lower, + tick_upper, + whirlpool_account.tick_spacing, + ) + .await?; for start_tick in tick_arrays.iter() { let (tick_array_address, _) = get_tick_array_address(&whirlpool, *start_tick)?; @@ -253,22 +328,22 @@ pub async fn setup_te_position( let te_position_token_account = get_associated_token_address(&owner, &te_position_mint.pubkey()); - let open_position_ix = OpenPosition { + let open_position_ix = OpenPositionWithTokenExtensions { funder: ctx.signer.pubkey(), owner, position: position_pubkey, position_mint: te_position_mint.pubkey(), position_token_account: te_position_token_account, whirlpool, - token_program: TOKEN_2022_PROGRAM_ID, + token2022_program: TOKEN_2022_PROGRAM_ID, system_program: system_program::id(), associated_token_program: spl_associated_token_account::id(), - rent: RENT_PROGRAM_ID, + metadata_update_auth, } - .instruction(OpenPositionInstructionArgs { + .instruction(OpenPositionWithTokenExtensionsInstructionArgs { tick_lower_index: lower_tick_index, tick_upper_index: upper_tick_index, - position_bump, + with_token_metadata_extension: true, }); ctx.send_transaction_with_signers(vec![open_position_ix], vec![&te_position_mint])