test: add interfaces for the simulation test

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";

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,

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,15 +7,20 @@ 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,
     ext_contract, is_promise_success,
     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(),

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);

makefile

@@ -58,14 +58,16 @@ 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
 	@cp ./res/mock_fungible_token.wasm ./tests/compiled-contracts/mock_fungible_token.wasm
 	@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

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/
+    yarn test -m "root sets*"       # same thing using yarn instead of npm, see https://yarnpkg.com/