From 436973830c93af68c40cecfd985620fb695dc05a Mon Sep 17 00:00:00 2001
From: Coach Chuck <169060940+coachchucksol@users.noreply.github.com>
Date: Tue, 23 Jul 2024 10:37:51 -0600
Subject: [PATCH] Fix rebalance tests (#59)

Co-authored-by: Evan Batsell <ebatsell@gmail.com>
---
 tests/tests/steward/test_integration.rs | 42 +++++++++++++++++++------
 1 file changed, 33 insertions(+), 9 deletions(-)

diff --git a/tests/tests/steward/test_integration.rs b/tests/tests/steward/test_integration.rs
index 7cbb7d5a..c3cae7eb 100644
--- a/tests/tests/steward/test_integration.rs
+++ b/tests/tests/steward/test_integration.rs
@@ -947,15 +947,6 @@ async fn test_rebalance_increase() {
         &serialized_stake_pool_account(stake_pool_spl, std::mem::size_of::<StakePool>()).into(),
     );
 
-    let mut steward_state_account: StewardStateAccount =
-        fixture.load_and_deserialize(&fixture.steward_state).await;
-
-    steward_state_account.state.num_pool_validators += 1;
-    ctx.borrow_mut().set_account(
-        &fixture.steward_state,
-        &serialized_steward_state_account(steward_state_account).into(),
-    );
-
     let (stake_account_address, transient_stake_account_address, withdraw_authority) =
         fixture.stake_accounts_for_validator(vote_account).await;
 
@@ -1023,6 +1014,17 @@ async fn test_rebalance_increase() {
 
     fixture.submit_transaction_assert_success(tx).await;
 
+    let mut steward_state_account: StewardStateAccount =
+        fixture.load_and_deserialize(&fixture.steward_state).await;
+
+    // Force validator into the active set, don't wait for next cycle
+    steward_state_account.state.num_pool_validators += 1;
+    steward_state_account.state.validators_added -= 1;
+    ctx.borrow_mut().set_account(
+        &fixture.steward_state,
+        &serialized_steward_state_account(steward_state_account).into(),
+    );
+
     let reserve_before_rebalance = fixture.get_account(&fixture.stake_pool_meta.reserve).await;
 
     let tx = Transaction::new_signed_with_payer(
@@ -1231,6 +1233,17 @@ async fn test_rebalance_decrease() {
     );
     fixture.submit_transaction_assert_success(tx).await;
 
+    let mut steward_state_account: StewardStateAccount =
+        fixture.load_and_deserialize(&fixture.steward_state).await;
+
+    // Force validator into the active set, don't wait for next cycle
+    // steward_state_account.state.num_pool_validators += 1;
+    steward_state_account.state.validators_added -= 1;
+    ctx.borrow_mut().set_account(
+        &fixture.steward_state,
+        &serialized_steward_state_account(steward_state_account).into(),
+    );
+
     // Simulating stake deposit
     let stake_account_data = fixture.get_account(&stake_account_address).await;
 
@@ -1469,6 +1482,17 @@ async fn test_rebalance_other_cases() {
     );
     fixture.submit_transaction_assert_success(tx).await;
 
+    let mut steward_state_account: StewardStateAccount =
+        fixture.load_and_deserialize(&fixture.steward_state).await;
+
+    // Force validator into the active set, don't wait for next cycle
+    steward_state_account.state.num_pool_validators += 1;
+    steward_state_account.state.validators_added -= 1;
+    ctx.borrow_mut().set_account(
+        &fixture.steward_state,
+        &serialized_steward_state_account(steward_state_account).into(),
+    );
+
     let rebalance_ix = Instruction {
         program_id: jito_steward::id(),
         accounts: jito_steward::accounts::Rebalance {