From 10a7e948b19d828b35e3e90d46344c9d4b1243c0 Mon Sep 17 00:00:00 2001
From: Jon Cinque <joncinque@pm.me>
Date: Fri, 23 Jun 2023 21:30:40 +0200
Subject: [PATCH] Fixup implementation and tests

---
 bot/src/rpc_client_utils.rs | 21 +++++++++++----------
 bot/src/stake_pool_v0.rs    | 34 ++++++++++++++++++++--------------
 2 files changed, 31 insertions(+), 24 deletions(-)

diff --git a/bot/src/rpc_client_utils.rs b/bot/src/rpc_client_utils.rs
index c538efd8..00484cdd 100644
--- a/bot/src/rpc_client_utils.rs
+++ b/bot/src/rpc_client_utils.rs
@@ -285,23 +285,24 @@ pub fn send_and_confirm_transactions_with_spinner(
     Err("Max retries exceeded".into())
 }
 
-pub(crate) fn get_stake_with_fallback(
+pub(crate) fn get_active_and_inactive_stake(
     rpc_client: &RpcClient,
     stake_address: &Pubkey,
 ) -> Result<(u64, u64), Box<dyn error::Error>> {
     let stake_activation = rpc_client.get_stake_activation(*stake_address, None);
+    let stake_balance = rpc_client.get_balance(stake_address).map_err(|err| {
+        format!(
+            "Unable to get stake account balance: {}: {}",
+            stake_address, err
+        )
+    })?;
     if let Ok(stake_activation) = stake_activation {
-        Ok((stake_activation.active, stake_activation.inactive))
-    } else {
         Ok((
-            0,
-            rpc_client.get_balance(stake_address).map_err(|err| {
-                format!(
-                    "Unable to get stake account balance: {}: {}",
-                    stake_address, err
-                )
-            })?,
+            stake_activation.active,
+            stake_balance - stake_activation.active,
         ))
+    } else {
+        Ok((0, stake_balance))
     }
 }
 
diff --git a/bot/src/stake_pool_v0.rs b/bot/src/stake_pool_v0.rs
index 2d64d348..885ac4fa 100644
--- a/bot/src/stake_pool_v0.rs
+++ b/bot/src/stake_pool_v0.rs
@@ -2,7 +2,7 @@ use {
     crate::{
         generic_stake_pool::*,
         rpc_client_utils::{
-            get_all_stake_by_staker, get_stake_with_fallback,
+            get_active_and_inactive_stake, get_all_stake_by_staker,
             send_and_confirm_transactions_with_spinner, MultiClient,
         },
     },
@@ -655,9 +655,10 @@ where
             validator_stake.vote_address,
         );
 
-        let (active_balance, inactive_balance) = get_stake_with_fallback(client, &stake_address)?;
+        let (active_balance, inactive_balance) =
+            get_active_and_inactive_stake(client, &stake_address)?;
         let (transient_active_balance, transient_inactive_balance) =
-            get_stake_with_fallback(client, &transient_stake_address)?;
+            get_active_and_inactive_stake(client, &transient_stake_address)?;
 
         let list = if validator_stake.priority {
             &mut priority_stake
@@ -848,7 +849,6 @@ where
 
 #[cfg(test)]
 mod test {
-    use crate::lamports_to_sol;
     use {
         super::*,
         crate::{
@@ -1116,13 +1116,16 @@ mod test {
 
         // after the first epoch, validators 0 and 1 are at their target levels but validator 2
         // needs one more epoch for the additional bonus stake to arrive
-        for (validator, expected_sol_balance) in
-            validators
-                .iter()
-                .zip(&[lamports_to_sol(MIN_STAKE_ACCOUNT_BALANCE), 10., 110.])
-        {
+        for (validator, expected_sol_balance) in validators.iter().zip(&[
+            // this one needs to maintain a bit more, since it has the extra rents
+            MIN_STAKE_ACCOUNT_BALANCE + DELEGATION_RENT * 2,
+            sol_to_lamports(10.),
+            sol_to_lamports(110.),
+        ]) {
             assert_eq!(
-                sol_to_lamports(*expected_sol_balance),
+                // since two movements have been done, expect two extra stake
+                // rent-exemption amounts in the account
+                *expected_sol_balance,
                 validator_stake_balance(&client, stake_pool.authorized_staker.pubkey(), validator,),
                 "stake balance mismatch for validator {}, expected {}",
                 validator.identity,
@@ -1151,12 +1154,15 @@ mod test {
 
         // after the second epoch, validator 2 is now has all the bonus stake
         for (validator, expected_sol_balance) in validators.iter().zip(&[
-            lamports_to_sol(MIN_STAKE_ACCOUNT_BALANCE),
-            10.,
-            320. - lamports_to_sol(MIN_STAKE_ACCOUNT_BALANCE),
+            // this account needs to maintain the extra rent-exempt amounts
+            MIN_STAKE_ACCOUNT_BALANCE + DELEGATION_RENT * 2,
+            // no change here still
+            sol_to_lamports(10.),
+            // this account doesn't get the extra rent-exempt amounts
+            sol_to_lamports(320.) - MIN_STAKE_ACCOUNT_BALANCE - DELEGATION_RENT * 2,
         ]) {
             assert_eq!(
-                sol_to_lamports(*expected_sol_balance),
+                *expected_sol_balance,
                 validator_stake_balance(&client, stake_pool.authorized_staker.pubkey(), validator,),
                 "stake balance mismatch for validator {}",
                 validator.identity