From d17617552fbfb506230771ade10a0c32c0568b43 Mon Sep 17 00:00:00 2001 From: Evan Batsell Date: Thu, 19 Sep 2024 15:29:25 -0400 Subject: [PATCH] Tests for lamport balance default --- programs/steward/src/delegation.rs | 6 +- tests/tests/steward/test_algorithms.rs | 42 ++++++++++- tests/tests/steward/test_state_methods.rs | 85 ++++++++++++++++++++++- 3 files changed, 128 insertions(+), 5 deletions(-) diff --git a/programs/steward/src/delegation.rs b/programs/steward/src/delegation.rs index 27fa58d4..78c02886 100644 --- a/programs/steward/src/delegation.rs +++ b/programs/steward/src/delegation.rs @@ -241,7 +241,7 @@ impl UnstakeState { }) } - fn stake_deposit_unstake( + pub fn stake_deposit_unstake( &self, state: &StewardState, index: usize, @@ -276,7 +276,7 @@ impl UnstakeState { Ok(0) } - fn instant_unstake( + pub fn instant_unstake( &self, state: &StewardState, index: usize, @@ -303,7 +303,7 @@ impl UnstakeState { Ok(0) } - fn scoring_unstake(&self, current_lamports: u64, target_lamports: u64) -> Result { + pub fn scoring_unstake(&self, current_lamports: u64, target_lamports: u64) -> Result { // If there are additional lamports to unstake on this validator and the total unstaked lamports is below the cap, destake to the target if self.scoring_unstake_total < self.scoring_unstake_cap { let lamports_above_target = current_lamports diff --git a/tests/tests/steward/test_algorithms.rs b/tests/tests/steward/test_algorithms.rs index 824fc036..d200f33f 100644 --- a/tests/tests/steward/test_algorithms.rs +++ b/tests/tests/steward/test_algorithms.rs @@ -1,7 +1,7 @@ // Unit tests for scoring, instant unstake, and delegation methods use anchor_lang::AnchorSerialize; use jito_steward::{ - constants::{EPOCH_DEFAULT, SORTED_INDEX_DEFAULT}, + constants::{EPOCH_DEFAULT, LAMPORT_BALANCE_DEFAULT, SORTED_INDEX_DEFAULT}, delegation::{ decrease_stake_calculation, increase_stake_calculation, RebalanceType, UnstakeState, }, @@ -1596,3 +1596,43 @@ fn test_decrease_stake_calculation() { _ => false, }); } + +#[test] +fn test_decrease_stake_default_lamports() { + // Given internal lamport balance set to default, test that no changes happen when doing stake deposit unstake + + let mut state = StateMachineFixtures::default().state; + + state.validator_lamport_balances[0] = LAMPORT_BALANCE_DEFAULT; + + let mut unstake_state = UnstakeState { + stake_deposit_unstake_total: 0, + stake_deposit_unstake_cap: 1000 * LAMPORTS_PER_SOL, + ..Default::default() + }; + + let test_cases = vec![ + // current_lamports, target_lamports + (1500 * LAMPORTS_PER_SOL, 1000 * LAMPORTS_PER_SOL), + (3000 * LAMPORTS_PER_SOL, 500 * LAMPORTS_PER_SOL), + ]; + + for (current_lamports, target_lamports) in test_cases { + let result = unstake_state + .stake_deposit_unstake(&state, 0, current_lamports, target_lamports) + .unwrap(); + + assert_eq!(result, 0, "Expected 0 unstake lamports, but got {}", result); + } + + // Test when stake_deposit_unstake_total reaches stake_deposit_unstake_cap + unstake_state.stake_deposit_unstake_total = unstake_state.stake_deposit_unstake_cap; + let result = unstake_state + .stake_deposit_unstake(&state, 0, 2000 * LAMPORTS_PER_SOL, 1000 * LAMPORTS_PER_SOL) + .unwrap(); + assert_eq!( + result, 0, + "Expected 0 unstake lamports when cap is reached, but got {}", + result + ); +} diff --git a/tests/tests/steward/test_state_methods.rs b/tests/tests/steward/test_state_methods.rs index 7b049001..8c41251b 100644 --- a/tests/tests/steward/test_state_methods.rs +++ b/tests/tests/steward/test_state_methods.rs @@ -8,7 +8,7 @@ use anchor_lang::{error::Error, AnchorSerialize}; use jito_steward::{ - constants::{MAX_VALIDATORS, SORTED_INDEX_DEFAULT}, + constants::{LAMPORT_BALANCE_DEFAULT, MAX_VALIDATORS, SORTED_INDEX_DEFAULT}, delegation::RebalanceType, errors::StewardError, Delegation, StewardStateEnum, @@ -751,3 +751,86 @@ fn test_rebalance() { } } } + +#[test] +fn test_rebalance_default_lamports() { + let fixtures = StateMachineFixtures::default(); + let mut state = fixtures.state; + let mut validator_list = fixtures.validator_list.clone(); + + // Case 1: Lamports default, has transient stake + state.validator_lamport_balances[0] = LAMPORT_BALANCE_DEFAULT; + state.state_tag = StewardStateEnum::Rebalance; + state.delegations[0..3].copy_from_slice(&[ + Delegation::new(1, 1), + Delegation::default(), + Delegation::default(), + ]); + state.scores[0..3].copy_from_slice(&[1_000_000_000, 0, 0]); + state.sorted_score_indices[0..3].copy_from_slice(&[0, 1, 2]); + + validator_list[0].transient_stake_lamports = 1000.into(); + let validator_list_bigvec = BigVec { + data: &mut validator_list.try_to_vec().unwrap(), + }; + + let res = state.rebalance( + fixtures.current_epoch, + 0, + &validator_list_bigvec, + 3000 * LAMPORTS_PER_SOL, + 0, + u64::from(validator_list[0].active_stake_lamports), + 0, + 0, + &fixtures.config.parameters, + ); + + println!("{:?}", res); + assert!(res.is_ok()); + match res.unwrap() { + RebalanceType::None => {} + _ => panic!("Expected RebalanceType::Increase"), + } + assert_eq!(state.validator_lamport_balances[0], LAMPORT_BALANCE_DEFAULT); + + // Case 2: Lamports not default, no transient stake + let mut state = fixtures.state; + state.state_tag = StewardStateEnum::Rebalance; + state.delegations[0..3].copy_from_slice(&[ + Delegation::new(1, 1), + Delegation::default(), + Delegation::default(), + ]); + state.scores[0..3].copy_from_slice(&[1_000_000_000, 0, 0]); + state.sorted_score_indices[0..3].copy_from_slice(&[0, 1, 2]); + + state.validator_lamport_balances[0] = LAMPORT_BALANCE_DEFAULT; + validator_list[0].transient_stake_lamports = 0.into(); + let validator_list_bigvec = BigVec { + data: &mut validator_list.try_to_vec().unwrap(), + }; + + let res = state.rebalance( + fixtures.current_epoch, + 0, + &validator_list_bigvec, + 4000 * LAMPORTS_PER_SOL, + 1000 * LAMPORTS_PER_SOL, + u64::from(validator_list[0].active_stake_lamports), + 0, + 0, + &fixtures.config.parameters, + ); + + assert!(res.is_ok()); + println!("{:?}", res); + if let RebalanceType::Increase(increase_amount) = res.unwrap() { + assert_eq!( + state.validator_lamport_balances[0], + 1000 * LAMPORTS_PER_SOL + increase_amount + ); + } else { + panic!("Expected RebalanceType::Increase"); + } +}