Skip to content

Commit

Permalink
Add decrease liquidity test for rust sdk (#608)
Browse files Browse the repository at this point in the history 
* 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
  • Loading branch information
@pauldragonfly
pauldragonfly authored Jan 16, 2025
1 parent 848f54d commit ead6aed
Show file tree
Hide file tree
Showing 2 changed files with 298 additions and 12 deletions.
223 changes: 217 additions & 6 deletions rust-sdk/whirlpool/src/decrease_liquidity.rs
View file
Original file line number Diff line number Diff line change
@@ -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,
};
Expand All @@ -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},
Expand Down Expand Up @@ -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<u64, Box<dyn Error>> {
let account_data = rpc.get_account(&address).await?;

if account_data.owner == TOKEN_2022_PROGRAM_ID {
// Token-2022 + (possibly) extension
let state = StateWithExtensionsOwned::<TokenAccount2022>::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<Position, Box<dyn Error>> {
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<dyn Error>> {
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<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 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(())
}
}
87 changes: 81 additions & 6 deletions rust-sdk/whirlpool/src/tests/program.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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<dyn Error>> {
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,
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -198,6 +264,7 @@ pub async fn setup_te_position(
tick_range: Option<(i32, i32)>,
owner: Option<Pubkey>,
) -> Result<Pubkey, Box<dyn Error>> {
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)?;
Expand All @@ -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)?;
Expand Down Expand Up @@ -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])
Expand Down

0 comments on commit ead6aed

Please sign in to comment.