diff --git a/contracts/linear/src/epoch_actions.rs b/contracts/linear/src/epoch_actions.rs
index 8c1b8a06..ed869677 100644
--- a/contracts/linear/src/epoch_actions.rs
+++ b/contracts/linear/src/epoch_actions.rs
@@ -1,5 +1,7 @@
 use crate::*;
-use near_sdk::{is_promise_success, log, near_bindgen, Balance};
+#[cfg(feature = "test")]
+use near_sdk::Promise;
+use near_sdk::{is_promise_success, log, near_bindgen, Balance, PromiseError, PromiseOrValue};
 
 use crate::errors::*;
 use crate::events::Event;
@@ -14,10 +16,108 @@ const MAX_SYNC_BALANCE_DIFF: Balance = 100;
 /// during each epoch.
 #[near_bindgen]
 impl LiquidStakingContract {
-    pub fn epoch_stake(&mut self) -> bool {
+    // `stake_to_validator` and `unstake_from_validator` are used to mock
+    // stake amounts and unstake amounts of validators at the beginnings
+    // of simulation tests
+    #[payable]
+    #[cfg(feature = "test")]
+    pub fn stake_to_validator(&mut self, validator_id: AccountId, amount: U128) -> Promise {
         self.assert_running();
         // make sure enough gas was given
-        let min_gas = GAS_EPOCH_STAKE + GAS_EXT_DEPOSIT_AND_STAKE + GAS_CB_VALIDATOR_STAKED;
+        let min_gas = GAS_EPOCH_STAKE
+            + GAS_EXT_DEPOSIT_AND_STAKE
+            + GAS_CB_VALIDATOR_STAKED
+            + GAS_SYNC_BALANCE
+            + GAS_CB_VALIDATOR_SYNC_BALANCE;
+
+        require!(
+            env::prepaid_gas() >= min_gas,
+            format!("{}. require at least {:?}", ERR_NO_ENOUGH_GAS, min_gas)
+        );
+
+        let mut validator = self
+            .validator_pool
+            .get_validator(&validator_id)
+            .expect(ERR_VALIDATOR_NOT_EXIST);
+
+        Event::EpochStakeAttempt {
+            validator_id: &validator_id,
+            amount: &amount,
+        }
+        .emit();
+
+        self.epoch_requested_stake_amount -= amount.0;
+
+        // do staking on selected validator
+        validator
+            .deposit_and_stake(&mut self.validator_pool, amount.into())
+            .then(ext_self_action_cb::validator_staked_callback(
+                validator.account_id.clone(),
+                amount.into(),
+                env::current_account_id(),
+                NO_DEPOSIT,
+                GAS_CB_VALIDATOR_STAKED + GAS_SYNC_BALANCE + GAS_CB_VALIDATOR_SYNC_BALANCE,
+            ))
+    }
+
+    #[cfg(feature = "test")]
+    pub fn unstake_from_validator(&mut self, validator_id: AccountId, amount: U128) -> Promise {
+        self.assert_running();
+        // make sure enough gas was given
+        let min_gas = GAS_EPOCH_UNSTAKE
+            + GAS_EXT_UNSTAKE
+            + GAS_CB_VALIDATOR_UNSTAKED
+            + GAS_SYNC_BALANCE
+            + GAS_CB_VALIDATOR_SYNC_BALANCE;
+        require!(
+            env::prepaid_gas() >= min_gas,
+            format!("{}. require at least {:?}", ERR_NO_ENOUGH_GAS, min_gas)
+        );
+
+        let mut validator = self
+            .validator_pool
+            .get_validator(&validator_id)
+            .expect(ERR_VALIDATOR_NOT_EXIST);
+
+        Event::EpochUnstakeAttempt {
+            validator_id: &validator_id,
+            amount: &amount,
+        }
+        .emit();
+
+        self.epoch_requested_unstake_amount -= amount.0;
+
+        // do staking on selected validator
+        validator
+            .unstake(&mut self.validator_pool, amount.into())
+            .then(ext_self_action_cb::validator_unstaked_callback(
+                validator.account_id.clone(),
+                amount.into(),
+                env::current_account_id(),
+                NO_DEPOSIT,
+                GAS_CB_VALIDATOR_STAKED + GAS_SYNC_BALANCE + GAS_CB_VALIDATOR_SYNC_BALANCE,
+            ))
+    }
+
+    /// Stake $NEAR to one of the validators.
+    ///
+    /// Select a candidate validator and stake part of or all of the to-settle
+    /// stake amounts to this validator. This function is expected to be called
+    /// in each epoch.
+    ///
+    /// # Return
+    /// * `true` - a candidate validator is selected and successfully staked to.
+    ///            There might be more stake amounts to settle so this function
+    ///            should be called again.
+    /// * `false` - There is no need to call this function again in this epoch.
+    pub fn epoch_stake(&mut self) -> PromiseOrValue<bool> {
+        self.assert_running();
+        // make sure enough gas was given
+        let min_gas = GAS_EPOCH_STAKE
+            + GAS_EXT_DEPOSIT_AND_STAKE
+            + GAS_CB_VALIDATOR_STAKED
+            + GAS_SYNC_BALANCE
+            + GAS_CB_VALIDATOR_SYNC_BALANCE;
         require!(
             env::prepaid_gas() >= min_gas,
             format!("{}. require at least {:?}", ERR_NO_ENOUGH_GAS, min_gas)
@@ -27,7 +127,7 @@ impl LiquidStakingContract {
         // after cleanup, there might be no need to stake
         if self.stake_amount_to_settle == 0 {
             log!("no need to stake, amount to settle is zero");
-            return false;
+            return PromiseOrValue::Value(false);
         }
 
         let candidate = self
@@ -36,7 +136,7 @@ impl LiquidStakingContract {
 
         if candidate.is_none() {
             log!("no candidate found to stake");
-            return false;
+            return PromiseOrValue::Value(false);
         }
 
         let mut candidate = candidate.unwrap();
@@ -44,7 +144,7 @@ impl LiquidStakingContract {
 
         if amount_to_stake < MIN_AMOUNT_TO_PERFORM_STAKE {
             log!("stake amount too low: {}", amount_to_stake);
-            return false;
+            return PromiseOrValue::Value(false);
         }
 
         require!(
@@ -70,16 +170,30 @@ impl LiquidStakingContract {
                 amount_to_stake.into(),
                 env::current_account_id(),
                 NO_DEPOSIT,
-                GAS_CB_VALIDATOR_STAKED,
-            ));
-
-        true
+                GAS_CB_VALIDATOR_STAKED + GAS_SYNC_BALANCE + GAS_CB_VALIDATOR_SYNC_BALANCE,
+            ))
+            .into()
     }
 
-    pub fn epoch_unstake(&mut self) -> bool {
+    /// Unstake $NEAR from one of the validators.
+    ///
+    /// Select a candidate validator and unstake part of or all of the to-settle
+    /// unstake amounts from this validator. This function is expected to be called
+    /// in each epoch.
+    ///
+    /// # Return
+    /// * `true` - a candidate validator is selected and successfully unstaked from.
+    ///            There might be more unstake amounts to settle so this function
+    ///            should be called again.
+    /// * `false` - There is no need to call this function again in this epoch.
+    pub fn epoch_unstake(&mut self) -> PromiseOrValue<bool> {
         self.assert_running();
         // make sure enough gas was given
-        let min_gas = GAS_EPOCH_UNSTAKE + GAS_EXT_UNSTAKE + GAS_CB_VALIDATOR_UNSTAKED;
+        let min_gas = GAS_EPOCH_UNSTAKE
+            + GAS_EXT_UNSTAKE
+            + GAS_CB_VALIDATOR_UNSTAKED
+            + GAS_SYNC_BALANCE
+            + GAS_CB_VALIDATOR_SYNC_BALANCE;
         require!(
             env::prepaid_gas() >= min_gas,
             format!("{}. require at least {:?}", ERR_NO_ENOUGH_GAS, min_gas)
@@ -89,7 +203,7 @@ impl LiquidStakingContract {
         // after cleanup, there might be no need to unstake
         if self.unstake_amount_to_settle == 0 {
             log!("no need to unstake, amount to settle is zero");
-            return false;
+            return PromiseOrValue::Value(false);
         }
 
         let candidate = self.validator_pool.get_candidate_to_unstake_v2(
@@ -98,14 +212,14 @@ impl LiquidStakingContract {
         );
         if candidate.is_none() {
             log!("no candidate found to unstake");
-            return false;
+            return PromiseOrValue::Value(false);
         }
         let mut candidate = candidate.unwrap();
         let amount_to_unstake = candidate.amount;
 
         if amount_to_unstake < MIN_AMOUNT_TO_PERFORM_UNSTAKE {
             log!("unstake amount too low: {}", amount_to_unstake);
-            return false;
+            return PromiseOrValue::Value(false);
         }
 
         // update internal state
@@ -126,10 +240,9 @@ impl LiquidStakingContract {
                 amount_to_unstake.into(),
                 env::current_account_id(),
                 NO_DEPOSIT,
-                GAS_CB_VALIDATOR_UNSTAKED,
-            ));
-
-        true
+                GAS_CB_VALIDATOR_UNSTAKED + GAS_SYNC_BALANCE + GAS_CB_VALIDATOR_SYNC_BALANCE,
+            ))
+            .into()
     }
 
     pub fn epoch_update_rewards(&mut self, validator_id: AccountId) {
@@ -225,32 +338,10 @@ impl LiquidStakingContract {
         .emit();
     }
 
-    /// Due to shares calculation and rounding of staking pool contract,
-    /// the amount of staked and unstaked balance might be a little bit
-    /// different than we requested.
-    /// This method is to sync the actual numbers with the validator.
-    pub fn sync_balance_from_validator(&mut self, validator_id: AccountId) {
-        self.assert_running();
-
-        let min_gas = GAS_SYNC_BALANCE + GAS_EXT_GET_ACCOUNT + GAS_CB_VALIDATOR_SYNC_BALANCE;
-        require!(
-            env::prepaid_gas() >= min_gas,
-            format!("{}. require at least {:?}", ERR_NO_ENOUGH_GAS, min_gas)
-        );
-
-        let mut validator = self
-            .validator_pool
-            .get_validator(&validator_id)
-            .expect(ERR_VALIDATOR_NOT_EXIST);
-
-        validator
-            .sync_account_balance(&mut self.validator_pool)
-            .then(ext_self_action_cb::validator_get_account_callback(
-                validator.account_id,
-                env::current_account_id(),
-                NO_DEPOSIT,
-                GAS_CB_VALIDATOR_SYNC_BALANCE,
-            ));
+    // To mock unsettled amounts at the beginnings of simulation tests
+    #[cfg(feature = "test")]
+    pub fn epoch_cleanup_for_test(&mut self) {
+        self.epoch_cleanup();
     }
 }
 
@@ -258,13 +349,21 @@ impl LiquidStakingContract {
 
 #[ext_contract(ext_self_action_cb)]
 trait EpochActionCallbacks {
-    fn validator_staked_callback(&mut self, validator_id: AccountId, amount: U128);
+    fn validator_staked_callback(
+        &mut self,
+        validator_id: AccountId,
+        amount: U128,
+    ) -> PromiseOrValue<bool>;
 
-    fn validator_unstaked_callback(&mut self, validator_id: AccountId, amount: U128);
+    fn validator_unstaked_callback(
+        &mut self,
+        validator_id: AccountId,
+        amount: U128,
+    ) -> PromiseOrValue<bool>;
 
     fn validator_get_balance_callback(&mut self, validator_id: AccountId);
 
-    fn validator_get_account_callback(&mut self, validator_id: AccountId);
+    fn validator_get_account_callback(&mut self, validator_id: AccountId) -> bool;
 
     fn validator_withdraw_callback(&mut self, validator_id: AccountId, amount: U128);
 }
@@ -273,8 +372,15 @@ trait EpochActionCallbacks {
 /// functions here SHOULD NOT PANIC!
 #[near_bindgen]
 impl LiquidStakingContract {
+    /// # Return
+    /// * `true` - Stake and sync balance succeed
+    /// * `false` - Stake fails
     #[private]
-    pub fn validator_staked_callback(&mut self, validator_id: AccountId, amount: U128) {
+    pub fn validator_staked_callback(
+        &mut self,
+        validator_id: AccountId,
+        amount: U128,
+    ) -> PromiseOrValue<bool> {
         let amount = amount.into();
         let mut validator = self
             .validator_pool
@@ -289,6 +395,16 @@ impl LiquidStakingContract {
                 amount: &U128(amount),
             }
             .emit();
+
+            validator
+                .sync_account_balance()
+                .then(ext_self_action_cb::validator_get_account_callback(
+                    validator_id,
+                    env::current_account_id(),
+                    NO_DEPOSIT,
+                    GAS_CB_VALIDATOR_SYNC_BALANCE,
+                ))
+                .into()
         } else {
             validator.on_stake_failed(&mut self.validator_pool);
 
@@ -300,11 +416,20 @@ impl LiquidStakingContract {
                 amount: &U128(amount),
             }
             .emit();
+
+            PromiseOrValue::Value(false)
         }
     }
 
+    /// # Return
+    /// * `true` - Unstake and sync balance succeed
+    /// * `false` - Unstake fails
     #[private]
-    pub fn validator_unstaked_callback(&mut self, validator_id: AccountId, amount: U128) {
+    pub fn validator_unstaked_callback(
+        &mut self,
+        validator_id: AccountId,
+        amount: U128,
+    ) -> PromiseOrValue<bool> {
         let amount = amount.into();
         let mut validator = self
             .validator_pool
@@ -319,6 +444,16 @@ impl LiquidStakingContract {
                 amount: &U128(amount),
             }
             .emit();
+
+            validator
+                .sync_account_balance()
+                .then(ext_self_action_cb::validator_get_account_callback(
+                    validator_id,
+                    env::current_account_id(),
+                    NO_DEPOSIT,
+                    GAS_CB_VALIDATOR_SYNC_BALANCE,
+                ))
+                .into()
         } else {
             // unstake failed, revert
             // 1. revert contract states
@@ -332,6 +467,8 @@ impl LiquidStakingContract {
                 amount: &U128(amount),
             }
             .emit();
+
+            PromiseOrValue::Value(false)
         }
     }
 
@@ -370,56 +507,71 @@ impl LiquidStakingContract {
     pub fn validator_get_account_callback(
         &mut self,
         validator_id: AccountId,
-        #[callback] account: HumanReadableAccount,
-    ) {
+        #[callback_result] result: Result<HumanReadableAccount, PromiseError>,
+    ) -> bool {
         let mut validator = self
             .validator_pool
             .get_validator(&validator_id)
             .unwrap_or_else(|| panic!("{}: {}", ERR_VALIDATOR_NOT_EXIST, &validator_id));
 
-        // allow at most MAX_SYNC_BALANCE_DIFF diff in total balance, staked balance and unstake balance
-        let new_total_balance = account.staked_balance.0 + account.unstaked_balance.0;
-        if abs_diff_eq(
-            new_total_balance,
-            validator.total_balance(),
-            MAX_SYNC_BALANCE_DIFF,
-        ) && abs_diff_eq(
-            account.staked_balance.0,
-            validator.staked_amount,
-            MAX_SYNC_BALANCE_DIFF,
-        ) && abs_diff_eq(
-            account.unstaked_balance.0,
-            validator.unstaked_amount,
-            MAX_SYNC_BALANCE_DIFF,
-        ) {
-            Event::SyncValidatorBalanceSuccess {
-                validator_id: &validator_id,
-                old_staked_balance: &validator.staked_amount.into(),
-                old_unstaked_balance: &validator.unstaked_amount.into(),
-                old_total_balance: &validator.total_balance().into(),
-                new_staked_balance: &account.staked_balance,
-                new_unstaked_balance: &account.unstaked_balance,
-                new_total_balance: &new_total_balance.into(),
+        match result {
+            Ok(account) => {
+                // allow at most MAX_SYNC_BALANCE_DIFF diff in total balance, staked balance and unstake balance
+                let new_total_balance = account.staked_balance.0 + account.unstaked_balance.0;
+                if abs_diff_eq(
+                    new_total_balance,
+                    validator.total_balance(),
+                    MAX_SYNC_BALANCE_DIFF,
+                ) && abs_diff_eq(
+                    account.staked_balance.0,
+                    validator.staked_amount,
+                    MAX_SYNC_BALANCE_DIFF,
+                ) && abs_diff_eq(
+                    account.unstaked_balance.0,
+                    validator.unstaked_amount,
+                    MAX_SYNC_BALANCE_DIFF,
+                ) {
+                    Event::SyncValidatorBalanceSuccess {
+                        validator_id: &validator_id,
+                        old_staked_balance: &validator.staked_amount.into(),
+                        old_unstaked_balance: &validator.unstaked_amount.into(),
+                        old_total_balance: &validator.total_balance().into(),
+                        new_staked_balance: &account.staked_balance,
+                        new_unstaked_balance: &account.unstaked_balance,
+                        new_total_balance: &new_total_balance.into(),
+                    }
+                    .emit();
+                    validator.on_sync_account_balance_success(
+                        &mut self.validator_pool,
+                        account.staked_balance.0,
+                        account.unstaked_balance.0,
+                    );
+                } else {
+                    Event::SyncValidatorBalanceFailedLargeDiff {
+                        validator_id: &validator_id,
+                        old_staked_balance: &validator.staked_amount.into(),
+                        old_unstaked_balance: &validator.unstaked_amount.into(),
+                        old_total_balance: &validator.total_balance().into(),
+                        new_staked_balance: &account.staked_balance,
+                        new_unstaked_balance: &account.unstaked_balance,
+                        new_total_balance: &new_total_balance.into(),
+                    }
+                    .emit();
+                    validator.on_sync_account_balance_failed(&mut self.validator_pool);
+                }
             }
-            .emit();
-            validator.on_sync_account_balance_success(
-                &mut self.validator_pool,
-                account.staked_balance.0,
-                account.unstaked_balance.0,
-            );
-        } else {
-            Event::SyncValidatorBalanceFailed {
-                validator_id: &validator_id,
-                old_staked_balance: &validator.staked_amount.into(),
-                old_unstaked_balance: &validator.unstaked_amount.into(),
-                old_total_balance: &validator.total_balance().into(),
-                new_staked_balance: &account.staked_balance,
-                new_unstaked_balance: &account.unstaked_balance,
-                new_total_balance: &new_total_balance.into(),
+            Err(_) => {
+                Event::SyncValidatorBalanceFailedCantGetAccount {
+                    validator_id: &validator_id,
+                    old_staked_balance: &validator.staked_amount.into(),
+                    old_unstaked_balance: &validator.unstaked_amount.into(),
+                    old_total_balance: &validator.total_balance().into(),
+                }
+                .emit();
+                validator.on_sync_account_balance_failed(&mut self.validator_pool);
             }
-            .emit();
-            validator.on_sync_account_balance_failed(&mut self.validator_pool);
-        }
+        };
+        true
     }
 
     #[private]
diff --git a/contracts/linear/src/errors.rs b/contracts/linear/src/errors.rs
index afcaa3e9..ecfddaeb 100644
--- a/contracts/linear/src/errors.rs
+++ b/contracts/linear/src/errors.rs
@@ -80,6 +80,8 @@ pub const ERR_VALIDATOR_UNSTAKE_WHEN_LOCKED: &str =
 pub const ERR_VALIDATOR_WITHDRAW_WHEN_LOCKED: &str =
     "Cannot withdraw from a pending release validator";
 pub const ERR_VALIDATOR_ALREADY_EXECUTING_ACTION: &str = "Validator is already executing action";
+pub const ERR_VALIDATOR_SYNC_BALANCE_NOT_ALLOWED: &str =
+    "Validator sync balance can only be called after stake or unstake";
 
 // liquidity pool
 pub const ERR_NON_POSITIVE_MIN_FEE: &str = "The min fee basis points should be positive";
diff --git a/contracts/linear/src/events.rs b/contracts/linear/src/events.rs
index e087d534..3c60fb6a 100644
--- a/contracts/linear/src/events.rs
+++ b/contracts/linear/src/events.rs
@@ -90,7 +90,7 @@ pub enum Event<'a> {
         new_unstaked_balance: &'a U128,
         new_total_balance: &'a U128,
     },
-    SyncValidatorBalanceFailed {
+    SyncValidatorBalanceFailedLargeDiff {
         validator_id: &'a AccountId,
         old_staked_balance: &'a U128,
         old_unstaked_balance: &'a U128,
@@ -99,6 +99,12 @@ pub enum Event<'a> {
         new_unstaked_balance: &'a U128,
         new_total_balance: &'a U128,
     },
+    SyncValidatorBalanceFailedCantGetAccount {
+        validator_id: &'a AccountId,
+        old_staked_balance: &'a U128,
+        old_unstaked_balance: &'a U128,
+        old_total_balance: &'a U128,
+    },
     // Staking Pool Interface
     Deposit {
         account_id: &'a AccountId,
diff --git a/contracts/linear/src/validator_pool.rs b/contracts/linear/src/validator_pool.rs
index 2116d1a3..5faf86ca 100644
--- a/contracts/linear/src/validator_pool.rs
+++ b/contracts/linear/src/validator_pool.rs
@@ -1,3 +1,4 @@
+use crate::epoch_actions::ext_self_action_cb;
 use crate::errors::*;
 use crate::events::Event;
 use crate::legacy::ValidatorV1_0_0;
@@ -6,6 +7,7 @@ use crate::legacy::ValidatorV1_4_0;
 use crate::types::*;
 use crate::utils::*;
 use crate::*;
+use near_sdk::PromiseOrValue;
 use near_sdk::{
     borsh::{self, BorshDeserialize, BorshSerialize},
     collections::UnorderedMap,
@@ -13,8 +15,12 @@ use near_sdk::{
     json_types::U128,
     near_bindgen, require, AccountId, Balance, EpochHeight, Promise,
 };
+
 use std::cmp::{max, min, Ordering};
 
+#[cfg(feature = "test")]
+use near_sdk::json_types::U64;
+
 const STAKE_SMALL_CHANGE_AMOUNT: Balance = ONE_NEAR;
 const UNSTAKE_FACTOR: u128 = 2;
 const MAX_UPDATE_WEIGHTS_COUNT: usize = 300;
@@ -160,6 +166,29 @@ impl ValidatorPool {
         old_weight
     }
 
+    // to mock pending release validators at the beginnings of simulation tests
+    #[cfg(feature = "test")]
+    pub fn set_unstake_fired_epoch(&mut self, validator_id: &AccountId, epoch_height: EpochHeight) {
+        let mut validator: Validator = self
+            .validators
+            .get(validator_id)
+            .expect(ERR_VALIDATOR_NOT_EXIST)
+            .into();
+        validator.unstake_fired_epoch = epoch_height;
+        self.validators.insert(validator_id, &validator.into());
+    }
+
+    // to mock draining validators at the beginnings of simulation tests
+    #[cfg(feature = "test")]
+    pub fn set_draining(&mut self, validator_id: &AccountId) {
+        let mut validator: Validator = self
+            .validators
+            .get(validator_id)
+            .expect(ERR_VALIDATOR_NOT_EXIST)
+            .into();
+        validator.set_draining(self, true);
+    }
+
     /// Update base stake amount of the validator
     pub fn update_base_stake_amount(&mut self, validator_id: &AccountId, amount: Balance) {
         let mut validator: Validator = self
@@ -584,6 +613,32 @@ impl LiquidStakingContract {
         }
     }
 
+    #[cfg(feature = "test")]
+    pub fn batch_set_unstake_fired_epoch(
+        &mut self,
+        validator_ids: Vec<AccountId>,
+        epoch_heights: Vec<U64>,
+    ) {
+        self.assert_running();
+        self.assert_manager();
+        for i in 0..validator_ids.len() {
+            self.validator_pool
+                .set_unstake_fired_epoch(&validator_ids[i], epoch_heights[i].into());
+        }
+    }
+
+    #[cfg(feature = "test")]
+    pub fn set_drainings(&mut self, validator_ids: Vec<AccountId>) {
+        for i in 0..validator_ids.len() {
+            self.validator_pool.set_draining(&validator_ids[i]);
+        }
+    }
+
+    #[cfg(feature = "test")]
+    pub fn set_total_staked_near_amount(&mut self, amount: U128) {
+        self.total_staked_near_amount = amount.0;
+    }
+
     // --- View Functions ---
 
     #[cfg(feature = "test")]
@@ -611,7 +666,11 @@ impl LiquidStakingContract {
 
 #[ext_contract(ext_self_validator_drain_cb)]
 trait ValidatorDrainCallbacks {
-    fn validator_drain_unstaked_callback(&mut self, validator_id: AccountId, amount: U128);
+    fn validator_drain_unstaked_callback(
+        &mut self,
+        validator_id: AccountId,
+        amount: U128,
+    ) -> PromiseOrValue<()>;
 
     fn validator_drain_withdraw_callback(&mut self, validator_id: AccountId, amount: U128);
 }
@@ -621,12 +680,16 @@ impl LiquidStakingContract {
     /// This method is designed to drain a validator.
     /// The weight of target validator should be set to 0 before calling this.
     /// And a following call to drain_withdraw MUST be made after 4 epoches.
-    pub fn drain_unstake(&mut self, validator_id: AccountId) {
+    pub fn drain_unstake(&mut self, validator_id: AccountId) -> Promise {
         self.assert_running();
         self.assert_manager();
 
         // make sure enough gas was given
-        let min_gas = GAS_DRAIN_UNSTAKE + GAS_EXT_UNSTAKE + GAS_CB_VALIDATOR_UNSTAKED;
+        let min_gas = GAS_DRAIN_UNSTAKE
+            + GAS_EXT_UNSTAKE
+            + GAS_CB_VALIDATOR_UNSTAKED
+            + GAS_SYNC_BALANCE
+            + GAS_CB_VALIDATOR_SYNC_BALANCE;
         require!(
             env::prepaid_gas() >= min_gas,
             format!("{}. require at least {:?}", ERR_NO_ENOUGH_GAS, min_gas)
@@ -676,9 +739,9 @@ impl LiquidStakingContract {
                     unstake_amount.into(),
                     env::current_account_id(),
                     NO_DEPOSIT,
-                    GAS_CB_VALIDATOR_UNSTAKED,
+                    GAS_CB_VALIDATOR_UNSTAKED + GAS_SYNC_BALANCE + GAS_CB_VALIDATOR_SYNC_BALANCE,
                 ),
-            );
+            )
     }
 
     /// Withdraw from a drained validator
@@ -735,7 +798,11 @@ impl LiquidStakingContract {
     }
 
     #[private]
-    pub fn validator_drain_unstaked_callback(&mut self, validator_id: AccountId, amount: U128) {
+    pub fn validator_drain_unstaked_callback(
+        &mut self,
+        validator_id: AccountId,
+        amount: U128,
+    ) -> PromiseOrValue<()> {
         let amount = amount.into();
         let mut validator = self
             .validator_pool
@@ -751,6 +818,16 @@ impl LiquidStakingContract {
                 amount: &U128(amount),
             }
             .emit();
+
+            validator
+                .sync_account_balance()
+                .then(ext_self_action_cb::validator_get_account_callback(
+                    validator_id,
+                    env::current_account_id(),
+                    NO_DEPOSIT,
+                    GAS_CB_VALIDATOR_SYNC_BALANCE,
+                ))
+                .into()
         } else {
             // unstake failed, revert
             validator.on_unstake_failed(&mut self.validator_pool);
@@ -760,6 +837,8 @@ impl LiquidStakingContract {
                 amount: &U128(amount),
             }
             .emit();
+
+            PromiseOrValue::Value(())
         }
     }
 
@@ -927,8 +1006,7 @@ impl Validator {
     }
 
     pub fn on_stake_success(&mut self, pool: &mut ValidatorPool, amount: Balance) {
-        self.post_execution(pool);
-
+        // Do not call post_execution() here because we need to sync account balance after stake
         self.staked_amount += amount;
         pool.save_validator(self);
     }
@@ -965,8 +1043,7 @@ impl Validator {
     }
 
     pub fn on_unstake_success(&mut self, pool: &mut ValidatorPool, amount: Balance) {
-        self.post_execution(pool);
-
+        // Do not call post_execution() here because we need to sync account balance after unstake
         self.staked_amount -= amount;
         self.unstaked_amount += amount;
         pool.save_validator(self);
@@ -1000,8 +1077,12 @@ impl Validator {
         pool.save_validator(self);
     }
 
-    pub fn sync_account_balance(&mut self, pool: &mut ValidatorPool) -> Promise {
-        self.pre_execution(pool);
+    /// Due to shares calculation and rounding of staking pool contract,
+    /// the amount of staked and unstaked balance might be a little bit
+    /// different than we requested.
+    /// This method is to sync the actual numbers with the validator.
+    pub fn sync_account_balance(&mut self) -> Promise {
+        require!(self.executing, ERR_VALIDATOR_SYNC_BALANCE_NOT_ALLOWED);
 
         ext_staking_pool::get_account(
             env::current_account_id(),
diff --git a/contracts/mock-staking-pool/src/lib.rs b/contracts/mock-staking-pool/src/lib.rs
index 43e7acb6..8662889a 100644
--- a/contracts/mock-staking-pool/src/lib.rs
+++ b/contracts/mock-staking-pool/src/lib.rs
@@ -49,6 +49,10 @@ pub struct MockStakingPool {
     staked: LookupMap<AccountId, u128>,
     /// for testing purpose, simulates contract panic
     panic: bool,
+    get_account_fail: bool,
+
+    staked_delta: u128,
+    unstaked_delta: u128,
 }
 
 #[near_bindgen]
@@ -59,6 +63,9 @@ impl MockStakingPool {
             deposits: LookupMap::new(b"d"),
             staked: LookupMap::new(b"s"),
             panic: false,
+            get_account_fail: false,
+            staked_delta: 0,
+            unstaked_delta: 0,
         }
     }
 }
@@ -84,6 +91,10 @@ impl StakingPool for MockStakingPool {
 
     fn get_account(&self, account_id: AccountId) -> HumanReadableAccount {
         require!(!self.panic, "Test Panic!");
+        require!(
+            !self.get_account_fail,
+            "get_account() failed, for testing purpose",
+        );
         HumanReadableAccount {
             account_id: account_id.clone(),
             staked_balance: U128::from(self.internal_get_staked(&account_id)),
@@ -157,17 +168,13 @@ impl MockStakingPool {
         self.panic = panic;
     }
 
-    pub fn adjust_balance(
-        &mut self,
-        account_id: AccountId,
-        staked_delta: U128,
-        unstaked_delta: U128,
-    ) {
-        let staked_amount = self.internal_get_staked(&account_id) - staked_delta.0;
-        let unstaked_amount = self.internal_get_unstaked_deposit(&account_id) + unstaked_delta.0;
+    pub fn set_get_account_fail(&mut self, value: bool) {
+        self.get_account_fail = value;
+    }
 
-        self.staked.insert(&account_id, &staked_amount);
-        self.deposits.insert(&account_id, &unstaked_amount);
+    pub fn set_balance_delta(&mut self, staked_delta: U128, unstaked_delta: U128) {
+        self.staked_delta = staked_delta.0;
+        self.unstaked_delta = unstaked_delta.0;
     }
 }
 
@@ -190,7 +197,7 @@ impl MockStakingPool {
         assert!(unstaked_deposit >= amount);
 
         let new_deposit = unstaked_deposit - amount;
-        let new_staked = self.internal_get_staked(&account_id) + amount;
+        let new_staked = self.internal_get_staked(&account_id) + amount - self.staked_delta;
 
         self.deposits.insert(&account_id, &new_deposit);
         self.staked.insert(&account_id, &new_staked);
@@ -202,7 +209,7 @@ impl MockStakingPool {
         assert!(staked >= amount);
 
         let unstaked_deposit = self.internal_get_unstaked_deposit(&account_id);
-        let new_deposit = unstaked_deposit + amount;
+        let new_deposit = unstaked_deposit + amount + self.unstaked_delta;
         let new_staked = staked - amount;
 
         self.deposits.insert(&account_id, &new_deposit);
diff --git a/makefile b/makefile
index ff5aa79d..6bb80652 100644
--- a/makefile
+++ b/makefile
@@ -58,7 +58,7 @@ test-unit:
 
 TEST_FILE ?= **
 LOGS ?=
-test-linear: linear_test mock-staking-pool mock-fungible-token mock-dex mock-lockup mock-whitelist
+test-contracts: linear_test mock-staking-pool mock-fungible-token mock-dex mock-lockup mock-whitelist
 	@mkdir -p ./tests/compiled-contracts/
 	@cp ./res/linear_test.wasm ./tests/compiled-contracts/linear.wasm
 	@cp ./res/mock_staking_pool.wasm ./tests/compiled-contracts/mock_staking_pool.wasm
@@ -66,6 +66,8 @@ test-linear: linear_test mock-staking-pool mock-fungible-token mock-dex mock-loc
 	@cp ./res/mock_dex.wasm ./tests/compiled-contracts/mock_dex.wasm
 	@cp ./res/mock_lockup.wasm ./tests/compiled-contracts/mock_lockup.wasm
 	@cp ./res/mock_whitelist.wasm ./tests/compiled-contracts/mock_whitelist.wasm
+
+test-linear: test-contracts
 	cd tests && NEAR_PRINT_LOGS=$(LOGS) npx near-workspaces-ava --timeout=2m __tests__/linear/$(TEST_FILE).ava.ts --verbose
 
 test-mock-staking-pool: mock-staking-pool
diff --git a/tests/README.md b/tests/README.md
index 7aadfacb..26576e89 100644
--- a/tests/README.md
+++ b/tests/README.md
@@ -17,4 +17,4 @@ To run only one test file:
 To run only one test:
 
     npm run test -- -m "root sets*" # matches tests with titles starting with "root sets"
-    yarn test -m "root sets*"       # same thing using yarn instead of npm, see https://yarnpkg.com/
\ No newline at end of file
+    yarn test -m "root sets*"       # same thing using yarn instead of npm, see https://yarnpkg.com/
diff --git a/tests/__tests__/linear/drain.ava.ts b/tests/__tests__/linear/drain.ava.ts
index e99c5199..76405468 100644
--- a/tests/__tests__/linear/drain.ava.ts
+++ b/tests/__tests__/linear/drain.ava.ts
@@ -6,7 +6,10 @@ import {
   setManager,
   assertValidatorAmountHelper,
   updateBaseStakeAmounts,
-  getValidator
+  getValidator,
+  epochUnstake,
+  epochStake,
+  assertHasLog
 } from "./helper";
 
 const workspace = initWorkSpace();
@@ -14,14 +17,7 @@ const workspace = initWorkSpace();
 async function stakeAll (signer: NearAccount, contract: NearAccount) {
   let run = true;
   while (run) {
-    run = await signer.call(
-      contract,
-      'epoch_stake',
-      {},
-      {
-        gas: Gas.parse('200 Tgas')
-      }
-    );
+    run = await epochStake(signer, contract);
   }
 }
 
@@ -33,6 +29,9 @@ workspace.test('Non-manager call drain methods', async (test, {contract, alice})
             'drain_unstake',
             {
                 validator_id: 'foo'
+            },
+            {
+              gas: "275 Tgas"
             }
         ),
         'Only manager can perform this action'
@@ -80,14 +79,7 @@ workspace.test('drain constraints', async (test, {contract, root, owner, alice,
   );
 
   // run stake
-  await bob.call(
-    contract,
-    'epoch_stake',
-    {},
-    {
-      gas: Gas.parse('200 Tgas')
-    }
-  );
+  await epochStake(bob, contract);
 
   // 1. cannot drain unstake when weight > 0
   await assertFailure(
@@ -99,7 +91,7 @@ workspace.test('drain constraints', async (test, {contract, root, owner, alice,
         validator_id: v1.accountId
       },
       {
-        gas: Gas.parse('200 Tgas')
+        gas: Gas.parse('275 Tgas')
       }
     ),
     'Validator weight must be zero for drain operation'
@@ -125,7 +117,7 @@ workspace.test('drain constraints', async (test, {contract, root, owner, alice,
         validator_id: v1.accountId
       },
       {
-        gas: Gas.parse('200 Tgas')
+        gas: Gas.parse('275 Tgas')
       }
     ),
     'Validator base stake amount must be zero for drain operation'
@@ -157,14 +149,7 @@ workspace.test('drain constraints', async (test, {contract, root, owner, alice,
     {}
   );
 
-  await bob.call(
-    contract,
-    'epoch_unstake',
-    {},
-    {
-      gas: Gas.parse('200 Tgas')
-    }
-  );
+  await epochUnstake(bob, contract);
 
   // validator now have unstaked balance > 0
   const assertValidator = assertValidatorAmountHelper(test, contract, owner);
@@ -180,7 +165,7 @@ workspace.test('drain constraints', async (test, {contract, root, owner, alice,
         validator_id: v1.accountId
       },
       {
-        gas: Gas.parse('200 Tgas')
+        gas: Gas.parse('275 Tgas')
       }
     ),
     'Cannot unstake from a pending release validator'
@@ -203,7 +188,7 @@ workspace.test('drain constraints', async (test, {contract, root, owner, alice,
         validator_id: v1.accountId
       },
       {
-        gas: Gas.parse('200 Tgas')
+        gas: Gas.parse('275 Tgas')
       }
     ),
     'Validator unstaked amount too large for drain unstake'
@@ -302,7 +287,7 @@ workspace.test('drain unstake and withdraw', async (test, {contract, root, owner
       validator_id: v1.accountId
     },
     {
-      gas: Gas.parse('200 Tgas')
+      gas: Gas.parse('275 Tgas')
     }
   );
 
@@ -356,3 +341,110 @@ workspace.test('drain unstake and withdraw', async (test, {contract, root, owner
   await assertValidator(v1, '0', '0');
   await assertValidator(v2, '60', '0');
 });
+
+workspace.test('drain unstake: get_account fails', async (test, {contract, root, owner, alice, bob}) => {
+  const manager = alice;
+  await setManager(root, contract, owner, manager);
+
+  const v1 = await createStakingPool(root, 'v1');
+  const v2 = await createStakingPool(root, 'v2');
+
+  // add validator
+  await manager.call(
+    contract,
+    'add_validator',
+    {
+      validator_id: v1.accountId,
+      weight: 10
+    },
+    {
+      gas: Gas.parse('100 Tgas')
+    }
+  );
+  await manager.call(
+    contract,
+    'add_validator',
+    {
+      validator_id: v2.accountId,
+      weight: 10
+    },
+    {
+      gas: Gas.parse('100 Tgas')
+    }
+  );
+
+  // update base stake amount of v1 to 20 NEAR
+  await updateBaseStakeAmounts(
+    contract,
+    manager,
+    [
+      v1.accountId,
+    ],
+    [
+      NEAR.parse("20")
+    ]
+  );
+
+  // user stake
+  await alice.call(
+    contract,
+    'deposit_and_stake',
+    {},
+    {
+      attachedDeposit: NEAR.parse('50')
+    }
+  );
+
+  // run stake
+  await stakeAll(bob, contract);
+
+  /**
+   * Steps to drain a validator
+   * 1. set weight to 0
+   * 2. set base stake amount to 0
+   * 3. call drain_unstake
+   * 4. call drain_withdraw
+   */
+
+  await manager.call(
+    contract,
+    'update_weight',
+    {
+      validator_id: v1.accountId,
+      weight: 0
+    }
+  );
+
+  // reset base stake amount to 0 NEAR
+  await updateBaseStakeAmounts(
+    contract,
+    manager,
+    [
+      v1.accountId,
+    ],
+    [
+      NEAR.parse("0")
+    ]
+  );
+
+  v1.call(
+    v1,
+    'set_get_account_fail',
+    {
+      value: true
+    }
+  );
+
+  const ret = await manager.call_raw(
+    contract,
+    'drain_unstake',
+    {
+      validator_id: v1.accountId
+    },
+    {
+      gas: Gas.parse('275 Tgas')
+    }
+  );
+
+  assertHasLog(test, ret, 'sync_validator_balance_failed_cant_get_account');
+});
diff --git a/tests/__tests__/linear/epoch-action-failure.ava.ts b/tests/__tests__/linear/epoch-action-failure.ava.ts
index e1cc2293..d8eff143 100644
--- a/tests/__tests__/linear/epoch-action-failure.ava.ts
+++ b/tests/__tests__/linear/epoch-action-failure.ava.ts
@@ -1,5 +1,5 @@
 import { NearAccount, NEAR, Gas } from "near-workspaces-ava";
-import { assertFailure, initWorkSpace, createStakingPool, getValidator } from "./helper";
+import { initWorkSpace, createStakingPool, getValidator, epochStake, epochUnstake, epochUnstakeCallRaw, epochStakeCallRaw, assertHasLog } from "./helper";
 
 const workspace = initWorkSpace();
 
@@ -37,7 +37,7 @@ function assertValidatorHelper(
   }
 }
 
-workspace.test('epoch stake failure', async (test, { root, contract, owner, alice }) => {
+workspace.test('epoch stake failure: deposit_and_stake fails', async (test, { root, contract, owner, alice }) => {
   const assertValidator = assertValidatorHelper(test, contract, owner);
 
   const v1 = await createStakingPool(root, 'v1');
@@ -66,20 +66,62 @@ workspace.test('epoch stake failure', async (test, { root, contract, owner, alic
 
   await setPanic(v1);
 
+  const ret = await epochStakeCallRaw(owner, contract);
+
+  test.is(ret.parseResult(), false);
+
+  assertHasLog(test, ret, 'epoch_stake_failed');
+
+  // nothing should be staked
+  await assertValidator(v1, '0', '0');
+});
+
+workspace.test('epoch stake failure: get_account fails', async (test, { root, contract, owner, alice }) => {
+  const assertValidator = assertValidatorHelper(test, contract, owner);
+
+  const v1 = await createStakingPool(root, 'v1');
+
   await owner.call(
     contract,
-    'epoch_stake',
+    'add_validator',
+    {
+      validator_id: v1.accountId,
+      weight: 10
+    },
+    {
+      gas: Gas.parse('100 Tgas')
+    }
+  );
+
+  // user stake
+  await alice.call(
+    contract,
+    'deposit_and_stake',
     {},
     {
-      gas: Gas.parse('200 Tgas')
+      attachedDeposit: NEAR.parse('50')
     }
   );
 
-  // nothing should be staked
-  await assertValidator(v1, '0', '0');
+  v1.call(
+    v1,
+    'set_get_account_fail',
+    {
+      value: true
+    }
+  );
+
+  const ret = await epochStakeCallRaw(owner, contract);
+
+  test.is(ret.parseResult(), true);
+
+  assertHasLog(test, ret, 'sync_validator_balance_failed_cant_get_account');
+
+  // stake still succeeded
+  await assertValidator(v1, '60', '0');
 });
 
-workspace.test('unstake failure', async (test, { root, contract, owner, alice }) => {
+workspace.test('epoch stake failure: balance diff too large', async (test, { root, contract, owner, alice }) => {
   const assertValidator = assertValidatorHelper(test, contract, owner);
 
   const v1 = await createStakingPool(root, 'v1');
@@ -106,17 +148,65 @@ workspace.test('unstake failure', async (test, { root, contract, owner, alice })
     }
   );
 
+  const MAX_SYNC_BALANCE_DIFF = NEAR.from(100);
+  const diff = MAX_SYNC_BALANCE_DIFF.addn(1);
+
   await owner.call(
+    v1,
+    'set_balance_delta',
+    {
+      staked_delta: diff.toString(10),
+      unstaked_delta: diff.toString(10),
+    },
+  );
+
+  const ret = await epochStakeCallRaw(owner, contract);
+
+  test.is(ret.parseResult(), true);
+
+  assertHasLog(test, ret, 'sync_validator_balance_failed_large_diff');
+
+  // stake still succeeded
+  await assertValidator(v1, '60', '0');
+});
+
+workspace.test('epoch unstake failure: unstake fails', async (test, { root, contract, owner, alice }) => {
+  const assertValidator = assertValidatorHelper(test, contract, owner);
+
+  const v1 = await createStakingPool(root, 'v1');
+
+  await owner.call(
+    contract,
+    'add_validator',
+    {
+      validator_id: v1.accountId,
+      weight: 10
+    },
+    {
+      gas: Gas.parse('100 Tgas')
+    }
+  );
+
+  // user stake
+  await alice.call(
     contract,
-    'epoch_stake',
+    'deposit_and_stake',
     {},
     {
-      gas: Gas.parse('200 Tgas')
+      attachedDeposit: NEAR.parse('50')
     }
   );
 
+  await epochStake(owner, contract);
+
   await assertValidator(v1, '60', '0');
 
+  await owner.call(
+    contract,
+    'set_epoch_height',
+    { epoch: 11 }
+  );
+
   // user unstake
   await alice.call(
     contract,
@@ -126,20 +216,79 @@ workspace.test('unstake failure', async (test, { root, contract, owner, alice })
 
   await setPanic(v1);
 
+  const ret = await epochUnstakeCallRaw(owner, contract);
+
+  test.is(ret.parseResult(), false);
+
+  assertHasLog(test, ret, 'epoch_unstake_failed');
+
+  // no unstake should actual happen
+  await assertValidator(v1, '60', '0');
+});
+
+workspace.test('epoch unstake failure: get_account fails', async (test, { root, contract, owner, alice }) => {
+  const assertValidator = assertValidatorHelper(test, contract, owner);
+
+  const v1 = await createStakingPool(root, 'v1');
+
   await owner.call(
     contract,
-    'epoch_unstake',
+    'add_validator',
+    {
+      validator_id: v1.accountId,
+      weight: 10
+    },
+    {
+      gas: Gas.parse('100 Tgas')
+    }
+  );
+
+  // user stake
+  await alice.call(
+    contract,
+    'deposit_and_stake',
     {},
     {
-      gas: Gas.parse('200 Tgas')
+      attachedDeposit: NEAR.parse('50')
     }
   );
 
-  // no unstake should actual happen
+  await epochStake(owner, contract);
+
   await assertValidator(v1, '60', '0');
+
+  await owner.call(
+    contract,
+    'set_epoch_height',
+    { epoch: 11 }
+  );
+
+  // user unstake
+  await alice.call(
+    contract,
+    'unstake',
+    { amount: NEAR.parse('10') }
+  );
+
+  v1.call(
+    v1,
+    'set_get_account_fail',
+    {
+      value: true
+    }
+  );
+
+  const ret = await epochUnstakeCallRaw(owner, contract);
+
+  test.is(ret.parseResult(), true);
+
+  assertHasLog(test, ret, 'sync_validator_balance_failed_cant_get_account');
+
+  // unstake still succeeded
+  await assertValidator(v1, '50', '10');
 });
 
-workspace.test('withdraw failure', async (test, { root, contract, owner, alice }) => {
+workspace.test('epoch unstake failure: balance diff too large', async (test, { root, contract, owner, alice }) => {
   const assertValidator = assertValidatorHelper(test, contract, owner);
 
   const v1 = await createStakingPool(root, 'v1');
@@ -166,15 +315,74 @@ workspace.test('withdraw failure', async (test, { root, contract, owner, alice }
     }
   );
 
+  await epochStake(owner, contract);
+
+  await assertValidator(v1, '60', '0');
+
+  await owner.call(
+    contract,
+    'set_epoch_height',
+    { epoch: 11 }
+  );
+
+  // user unstake
+  await alice.call(
+    contract,
+    'unstake',
+    { amount: NEAR.parse('10') }
+  );
+
+  const MAX_SYNC_BALANCE_DIFF = NEAR.from(100);
+  const diff = MAX_SYNC_BALANCE_DIFF.addn(1);
+
+  await owner.call(
+    v1,
+    'set_balance_delta',
+    {
+      staked_delta: diff.toString(10),
+      unstaked_delta: diff.toString(10),
+    },
+  );
+
+  const ret = await epochUnstakeCallRaw(owner, contract);
+
+  test.is(ret.parseResult(), true);
+
+  assertHasLog(test, ret, 'sync_validator_balance_failed_large_diff');
+
+  // unstake still succeeded
+  await assertValidator(v1, '50', '10');
+});
+
+workspace.test('withdraw failure', async (test, { root, contract, owner, alice }) => {
+  const assertValidator = assertValidatorHelper(test, contract, owner);
+
+  const v1 = await createStakingPool(root, 'v1');
+
   await owner.call(
     contract,
-    'epoch_stake',
+    'add_validator',
+    {
+      validator_id: v1.accountId,
+      weight: 10
+    },
+    {
+      gas: Gas.parse('100 Tgas')
+    }
+  );
+
+  // user stake
+  await alice.call(
+    contract,
+    'deposit_and_stake',
     {},
     {
-      gas: Gas.parse('200 Tgas')
+      attachedDeposit: NEAR.parse('50')
     }
   );
 
+  await epochStake(owner, contract);
+
   // fast-forward 4 epoch
   await owner.call(
     contract,
@@ -191,14 +399,7 @@ workspace.test('withdraw failure', async (test, { root, contract, owner, alice }
     { amount: NEAR.parse('10') }
   );
 
-  await owner.call(
-    contract,
-    'epoch_unstake',
-    {},
-    {
-      gas: Gas.parse('200 Tgas')
-    }
-  );
+  await epochUnstake(owner, contract);
 
   await assertValidator(v1, '50', '10');
 
@@ -254,14 +455,7 @@ workspace.test('get balance failure', async (test, { root, contract, owner, alic
     }
   );
 
-  await owner.call(
-    contract,
-    'epoch_stake',
-    {},
-    {
-      gas: Gas.parse('200 Tgas')
-    }
-  );
+  await epochStake(owner, contract);
 
   await assertValidator(v1, '60', '0');
 
diff --git a/tests/__tests__/linear/epoch-action.ava.ts b/tests/__tests__/linear/epoch-action.ava.ts
index 49db43df..566f5e32 100644
--- a/tests/__tests__/linear/epoch-action.ava.ts
+++ b/tests/__tests__/linear/epoch-action.ava.ts
@@ -6,8 +6,9 @@ import {
   updateBaseStakeAmounts,
   setManager,
   assertValidatorAmountHelper,
-  getSummary,
-  skip
+  skip,
+  epochStake,
+  epochUnstake
 } from "./helper";
 
 const workspace = initWorkSpace();
@@ -15,28 +16,14 @@ const workspace = initWorkSpace();
 async function stakeAll (owner: NearAccount, contract: NearAccount) {
   let run = true;
   while (run) {
-    run = await owner.call(
-      contract,
-      'epoch_stake',
-      {},
-      {
-        gas: Gas.parse('200 Tgas')
-      }
-    );
+    run = await epochStake(owner, contract);
   }
 }
 
 async function unstakeAll (owner: NearAccount, contract: NearAccount) {
   let run = true;
   while (run) {
-    run = await owner.call(
-      contract,
-      'epoch_unstake',
-      {},
-      {
-        gas: Gas.parse('200 Tgas')
-      }
-    );
+    run = await epochUnstake(owner, contract);
   }
 }
 
@@ -774,7 +761,7 @@ skip('estimate gas of epoch unstake', async (test, {contract, alice, root, owner
       'epoch_unstake',
       {},
       {
-        gas: Gas.parse('300 Tgas')
+        gas: Gas.parse('280 Tgas')
       }
     );
 
diff --git a/tests/__tests__/linear/helper.ts b/tests/__tests__/linear/helper.ts
index bb1fc8f8..47fdee21 100644
--- a/tests/__tests__/linear/helper.ts
+++ b/tests/__tests__/linear/helper.ts
@@ -1,4 +1,4 @@
-import { Workspace, NEAR, NearAccount, BN } from "near-workspaces-ava";
+import { Workspace, NEAR, NearAccount, BN, Gas, TransactionResult } from "near-workspaces-ava";
 
 export const ONE_YOCTO = '1';
 export const NUM_EPOCHS_TO_UNLOCK = 4;
@@ -334,3 +334,63 @@ export function assertValidatorAmountHelper (
     }
   }
 }
+
+const EPOCH_STAKE_AND_UNSTAKE_GAS = Gas.parse('280 Tgas');
+
+export function epochStake(caller: NearAccount, contract: NearAccount): Promise<any> {
+  return caller.call(
+    contract,
+    'epoch_stake',
+    {},
+    {
+      gas: EPOCH_STAKE_AND_UNSTAKE_GAS
+    }
+  );
+}
+
+export function epochStakeCallRaw(caller: NearAccount, contract: NearAccount): Promise<any> {
+  return caller.call_raw(
+    contract,
+    'epoch_stake',
+    {},
+    {
+      gas: EPOCH_STAKE_AND_UNSTAKE_GAS
+    }
+  );
+}
+
+export function epochUnstake(caller: NearAccount, contract: NearAccount): Promise<any> {
+  return caller.call(
+    contract,
+    'epoch_unstake',
+    {},
+    {
+      gas: EPOCH_STAKE_AND_UNSTAKE_GAS
+    }
+  );
+}
+
+export function epochUnstakeCallRaw(caller: NearAccount, contract: NearAccount): Promise<any> {
+  return caller.call_raw(
+    contract,
+    'epoch_unstake',
+    {},
+    {
+      gas: EPOCH_STAKE_AND_UNSTAKE_GAS
+    }
+  );
+}
+
+export function assertHasLog(
+  test: any,
+  txResult: TransactionResult,
+  expected: string,
+) {
+  test.truthy(
+    txResult.result.receipts_outcome.find(
+      (outcome: any) => outcome.outcome.logs.find(
+        (log: any) => log.includes(expected)
+      )
+    )
+  );
+}
diff --git a/tests/__tests__/linear/staking-pool-interface.ava.ts b/tests/__tests__/linear/staking-pool-interface.ava.ts
index 4bffd93f..da3231c3 100644
--- a/tests/__tests__/linear/staking-pool-interface.ava.ts
+++ b/tests/__tests__/linear/staking-pool-interface.ava.ts
@@ -1,5 +1,5 @@
 import { NEAR, Gas } from 'near-workspaces-ava';
-import { initWorkSpace, assertFailure, epochHeightFastforward } from './helper';
+import { initWorkSpace, assertFailure, epochHeightFastforward, epochStake } from './helper';
 
 const ERR_UNSTAKED_BALANCE_NOT_AVAILABLE = 'The unstaked balance is not yet available due to unstaking delay';
 
@@ -268,14 +268,7 @@ workspace.test('late unstake and withdraw', async (test, { contract ,alice }) =>
   );
 
   // call epoch_stake, in order to trigger clean up
-  await alice.call(
-    contract,
-    'epoch_stake',
-    {},
-    {
-      gas: Gas.parse('200 Tgas')
-    }
-  );
+  await epochStake(alice, contract);
 
   // unstake
   const unstakeAmount = NEAR.parse('5');
diff --git a/tests/__tests__/linear/sync_balance.ava.ts b/tests/__tests__/linear/sync_balance.ava.ts
index 9a1e16dc..79dbc51e 100644
--- a/tests/__tests__/linear/sync_balance.ava.ts
+++ b/tests/__tests__/linear/sync_balance.ava.ts
@@ -1,4 +1,4 @@
-import { assertFailure, createStakingPool, getValidator, initWorkSpace } from "./helper";
+import { assertFailure, createStakingPool, epochStake, epochUnstake, getValidator, initWorkSpace } from "./helper";
 import { Gas, NEAR, NearAccount, ONE_NEAR, stake, } from "near-workspaces-ava";
 
 const MAX_SYNC_BALANCE_DIFF = NEAR.from(100);
@@ -67,67 +67,43 @@ workspace.test('sync balance failure', async (test, { root, contract, alice, own
     }
   );
 
-  for (let i = 0; i < 2; i++) {
-    await owner.call(
-      contract,
-      'epoch_stake',
-      {},
-      {
-        gas: Gas.parse('200 Tgas')
-      }
-    );
-  }
-
   // -- 1. total balance diff > MAX_SYNC_BALANCE_DIFF
   const diff = MAX_SYNC_BALANCE_DIFF.addn(1);
+
   await owner.call(
     v1,
-    'adjust_balance',
+    'set_balance_delta',
     {
-      account_id: contract.accountId,
-      staked_delta: "0",
-      unstaked_delta: diff.toString(10)
+      staked_delta: diff.toString(10),
+      unstaked_delta: diff.toString(10),
     },
   );
 
-  await owner.call(
-    contract,
-    'sync_balance_from_validator',
-    {
-      validator_id: v1.accountId
-    },
-    {
-      gas: Gas.parse('200 Tgas')
-    }
-  );
+  for (let i = 0; i < 2; i++) {
+    await epochStake(owner, contract);
+  }
 
   // v1 amount should not change
-  await assertValidator(v1, '30000000000000000000000000', '0');
+  await assertValidator(v1, NEAR.parse('30').toString(10), '0');
 
-  // -- 2. amount balance diff > MAX_SYNC_BALANCE_DIFF
   await owner.call(
-    v2,
-    'adjust_balance',
-    {
-      account_id: contract.accountId,
-      staked_delta: diff.toString(10),
-      unstaked_delta: diff.toString(10)
-    },
+    contract,
+    'set_epoch_height',
+    { epoch: 11 }
   );
 
-  await owner.call(
+  await alice.call(
     contract,
-    'sync_balance_from_validator',
-    {
-      validator_id: v2.accountId
-    },
-    {
-      gas: Gas.parse('200 Tgas')
-    }
+    'unstake_all',
+    {},
   );
 
+  for (let i = 0; i < 2; i++) {
+    await epochUnstake(owner, contract);
+  }
+
   // v2 amount should not change
-  await assertValidator(v2, '30000000000000000000000000', '0');
+  await assertValidator(v2, NEAR.parse('5').toString(10), NEAR.parse('25').toString(10));
 });
 
 workspace.test('sync balance', async (test, { root, contract, alice, owner }) => {
@@ -168,39 +144,47 @@ workspace.test('sync balance', async (test, { root, contract, alice, owner }) =>
     }
   );
 
+   // -- amount balance diff < MAX_SYNC_BALANCE_DIFF
+   const diff = MAX_SYNC_BALANCE_DIFF.subn(1);
+   await owner.call(
+     v2,
+     'set_balance_delta',
+     {
+       staked_delta: diff.toString(10),
+       unstaked_delta: diff.toString(10),
+     },
+   );
+
   for (let i = 0; i < 2; i++) {
-    await owner.call(
-      contract,
-      'epoch_stake',
-      {},
-      {
-        gas: Gas.parse('200 Tgas')
-      }
-    );
+    await epochStake(owner, contract);
   }
 
-  // -- amount balance diff < MAX_SYNC_BALANCE_DIFF
-  const diff = MAX_SYNC_BALANCE_DIFF.subn(1);
+  await assertValidator(v2, NEAR.parse("30").sub(diff).toString(10), '0');
+
+  await owner.call(
+    contract,
+    'set_epoch_height',
+    { epoch: 11 }
+  );
+
   await owner.call(
-    v2,
-    'adjust_balance',
+    v1,
+    'set_balance_delta',
     {
-      account_id: contract.accountId,
       staked_delta: diff.toString(10),
       unstaked_delta: diff.toString(10),
     },
   );
 
-  await owner.call(
+  await alice.call(
     contract,
-    'sync_balance_from_validator',
-    {
-      validator_id: v2.accountId
-    },
-    {
-      gas: Gas.parse('200 Tgas')
-    }
+    'unstake_all',
+    {},
   );
 
-  await assertValidator(v2, NEAR.parse("30").sub(diff).toString(10), diff.toString(10));
+  for (let i = 0; i < 2; i++) {
+    await epochUnstake(owner, contract);
+  }
+
+  await assertValidator(v1, NEAR.parse("5").toString(10),  NEAR.parse("25").add(diff).toString(10));
 });
diff --git a/tests/__tests__/linear/upgrade.ava.ts b/tests/__tests__/linear/upgrade.ava.ts
index b4c41744..13618d1a 100644
--- a/tests/__tests__/linear/upgrade.ava.ts
+++ b/tests/__tests__/linear/upgrade.ava.ts
@@ -1,7 +1,7 @@
 import { readFileSync } from "fs";
 import { Gas, NEAR } from "near-units";
 import { NearAccount, Workspace } from "near-workspaces-ava";
-import { createStakingPool, getValidator, initAndSetWhitelist, skip, updateBaseStakeAmounts, } from "./helper";
+import { createStakingPool, epochStake, getValidator, initAndSetWhitelist, skip, updateBaseStakeAmounts, } from "./helper";
 
 async function deployLinearAtVersion(
   root: NearAccount,
@@ -34,14 +34,7 @@ async function upgrade(contract: NearAccount, owner: NearAccount) {
 async function stakeAll (signer: NearAccount, contract: NearAccount) {
   let run = true;
   while (run) {
-    run = await signer.call(
-      contract,
-      'epoch_stake',
-      {},
-      {
-        gas: Gas.parse('300 Tgas')
-      }
-    );
+    run = await epochStake(signer, contract);
   }
 }
 
@@ -300,7 +293,7 @@ skip('upgrade from v1.3.3 to v1.4.0', async (test, context) => {
       validator_id: targetValidator.accountId
     },
     {
-      gas: Gas.parse('200 Tgas')
+      gas: Gas.parse('275 Tgas')
     }
   );
 
@@ -392,14 +385,7 @@ skip('upgrade from v1.4.4 to v1.5.0', async (test, context) => {
 
   async function delayedEpochStake(ms: number) {
     await sleep(ms);
-    await alice.call(
-      contract,
-      'epoch_stake',
-      {},
-      {
-        gas: Gas.parse('300 Tgas')
-      }
-    );
+    await epochStake(alice, contract);
   }
 
   let executed = false;