First Integration Test: Staking

We didn't put everything in a single spec.rs file, allowing us to split test into multiple files and importing them into main.rs. If the test is long, we can have one test function per file. Let's create our first test in its own file. In the same directory as utils.rs, make a test_staking.rs.

Make sure you put the macro in the main.rs file for general use.

use lockup::{
  LockupContractContract, TerminationStatus, TransfersInformation, VestingSchedule,
  VestingScheduleOrHash, VestingScheduleWithSalt, WrappedBalance, MIN_BALANCE_FOR_STORAGE
};
use near_sdk::borsh::BorshSerialize;
use near_sdk::json_types::U128;
use near_sdk::serde_json::json;
use near_sdk::{AccountId, Balance};
use near_sdk_sim::runtime::GenesisConfig;
use near_sdk_sim::{deploy, init_simulator, to_yocto, UserAccount, STORAGE_AMOUNT};
use quickcheck_macros::quickcheck;
use std::convert::TryInto;

mod test_staking_with_helpers;
mod test_termination_with_staking_hashed;

pub(crate) mod utils;  // utils require pub-crate
mod test_staking;  // import other test files like this. 

pub(crate) use crate::utils::*;

pub(crate) fn assert_almost_eq_with_max_delta(left: u128, right: u128, max_delta: u128) {
  assert!(
      std::cmp::max(left, right) - std::cmp::min(left, right) <= max_delta,
      "{}",
      format!(
          "Left {} is not even close to Right {} within delta {}",
          left, right, max_delta
      )
  );
}

pub(crate) fn assert_eq_with_gas(left: u128, right: u128) {
  assert_almost_eq_with_max_delta(left, right, to_yocto("0.005"));
}

pub(crate) fn assert_yocto_eq(left: u128, right: u128) {
  assert_almost_eq_with_max_delta(left, right, 1);
}


near_sdk_sim::lazy_static_include::lazy_static_include_bytes! {
  LOCKUP_WASM_BYTES => "res/lockup.wasm",
  STAKING_POOL_WASM_BYTES => "tests/res/staking_pool.wasm",
  FAKE_VOTING_WASM_BYTES => "tests/res/fake_voting.wasm",
  WHITELIST_WASM_BYTES => "tests/res/whitelist.wasm"
}

Then, we shall define some functions in main.rs that are used in test functions later. Ensure they have pub(crate) so it could be shared across modules.

use lockup::{
  LockupContractContract, TerminationStatus, TransfersInformation, VestingSchedule,
  VestingScheduleOrHash, VestingScheduleWithSalt, WrappedBalance, MIN_BALANCE_FOR_STORAGE
};
use near_sdk::borsh::BorshSerialize;
use near_sdk::json_types::U128;
use near_sdk::serde_json::json;
use near_sdk::{AccountId, Balance};
use near_sdk_sim::runtime::GenesisConfig;
use near_sdk_sim::{deploy, init_simulator, to_yocto, UserAccount, STORAGE_AMOUNT};
use quickcheck_macros::quickcheck;
use std::convert::TryInto;

mod test_staking_with_helpers;
mod test_termination_with_staking_hashed;

pub(crate) mod utils;  // utils require pub-crate
mod test_staking;  // import other test files like this. 

pub(crate) use crate::utils::*;

pub(crate) fn assert_almost_eq_with_max_delta(left: u128, right: u128, max_delta: u128) {
  assert!(
      std::cmp::max(left, right) - std::cmp::min(left, right) <= max_delta,
      "{}",
      format!(
          "Left {} is not even close to Right {} within delta {}",
          left, right, max_delta
      )
  );
}

pub(crate) fn assert_eq_with_gas(left: u128, right: u128) {
  assert_almost_eq_with_max_delta(left, right, to_yocto("0.005"));
}

pub(crate) fn assert_yocto_eq(left: u128, right: u128) {
  assert_almost_eq_with_max_delta(left, right, 1);
}


near_sdk_sim::lazy_static_include::lazy_static_include_bytes! {
  LOCKUP_WASM_BYTES => "res/lockup.wasm",
  STAKING_POOL_WASM_BYTES => "tests/res/staking_pool.wasm",
  FAKE_VOTING_WASM_BYTES => "tests/res/fake_voting.wasm",
  WHITELIST_WASM_BYTES => "tests/res/whitelist.wasm"
}

The first thing we need to do is to deploy the contract.

use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

This uses the deploy! macro. The signer_account is root, which is not a good imitation and shows the limitation of simulation testing. In reality, some account will deploy the contract instead of the "genesis" doing it like what is did here.

Another init_method is quite easy to understand: the 8 parameters matches those parameters passed to new. Let's recall:

use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
// use near_sdk::json_types::Base58PublicKey;
use near_sdk::{AccountId, env, ext_contract, near_bindgen, require};
// use near_account_id::AccountId;

pub use crate::foundation::*;
pub use crate::foundation_callbacks::*;
pub use crate::getters::*;
pub use crate::internal::*;
pub use crate::owner::*;
pub use crate::owner_callbacks::*;
pub use crate::types::*;

pub mod foundation;
pub mod foundation_callbacks;
pub mod gas;
pub mod owner_callbacks;
pub mod types;

pub mod getters;
pub mod internal;
pub mod owner;

// Not needed as of 4.0.0-pre.1. 
// #[global_allocator]
// static ALLOC: near_sdk::wee_alloc::WeeAlloc = near_sdk::wee_alloc::WeeAlloc::INIT;

const NO_DEPOSIT: u128 = 0;

/// At least 3.5 NEAR to avoid being transferred out to cover 
/// contract code storage and some internal state. 
pub const MIN_BALANCE_FOR_STORAGE: u128 = 3_500_000_000_000_000_000_000_000;

#[ext_contract(ext_staking_pool)]
pub trait ExtStakingPool {
    fn get_account_staked_balance(&self, account_id: AccountId) -> WrappedBalance;
    fn get_account_unstaked_balance(&self, account_id: AccountId) -> WrappedBalance;
    fn get_account_total_balance(&self, account_id: AccountId) -> WrappedBalance;

    fn deposit(&mut self);
    fn deposit_and_stake(&mut self);

    fn withdraw(&mut self, amount: WrappedBalance);

    fn stake(&mut self, amount: WrappedBalance);
    fn unstake(&mut self, amount: WrappedBalance);

    fn unstake_all(&mut self);
}

#[ext_contract(ext_whitelist)]
pub trait ExtStakingPoolWhitelist {
    fn is_whitelisted(&self, staking_pool_account_id: AccountId) -> bool;
}

#[ext_contract(ext_transfer_poll)]
pub trait ExtTransferPoll {
    fn get_result(&self) -> Option<PollResult>;
}

#[ext_contract(ext_self_owner)]
pub trait ExtLockupContractOwner {
    fn on_whitelist_is_whitelisted(
      &mut self, 
      #[callback] is_whitelisted: bool,
      staking_pool_account_id: AccountId,
    ) -> bool;

    fn on_staking_pool_deposit(&mut self, amount: WrappedBalance) -> bool;
    fn on_staking_pool_deposit_and_stake(&mut self, amount: WrappedBalance) -> bool;

    fn on_staking_pool_withdraw(&mut self, amount: WrappedBalance) -> bool;
    
    fn on_staking_pool_stake(&mut self, amount: WrappedBalance) -> bool;
    fn on_staking_pool_unstake(&mut self, amount: WrappedBalance) -> bool;

    fn on_staking_pool_unstake_all(&mut self) -> bool;

    fn on_get_result_from_transfer_poll(
      &mut self,
      #[callback] poll_result: PollResult
    ) -> bool;

    fn on_get_account_total_balance(
      &mut self,
      #[callback] total_balance: WrappedBalance
    );

    // don't be confused the one "by foundation". 
    fn on_get_account_unstaked_balance_to_withdraw_by_owner(
      &mut self,
      #[callback] unstaked_balance: WrappedBalance,
    );
}

#[ext_contract(ext_self_foundation)]
pub trait ExtLockupContractFoundation {
    fn on_withdraw_unvested_amount(
      &mut self,
      amount: WrappedBalance,
      receiver_id: AccountId,
    ) -> bool;

    fn on_get_account_staked_balance_to_unstake(
      &mut self,
      #[callback] staked_balance: WrappedBalance,
    );

    fn on_staking_pool_unstake_for_termination(
      &mut self,
      amount: WrappedBalance
    ) -> bool;

    fn on_get_account_unstaked_balance_to_withdraw(
      &mut self,
      #[callback] unstaked_balance: WrappedBalance,
    );

    fn on_staking_pool_withdraw_for_termination(
      &mut self,
      amount: WrappedBalance
    ) -> bool;
}

#[near_bindgen]
#[derive(BorshDeserialize, BorshSerialize)]
pub struct LockupContract {
    pub owner_account_id: AccountId,
    pub lockup_information: LockupInformation,  // schedule and amount
    pub vesting_information: VestingInformation,  // schedule and termination status
    pub staking_pool_whitelist_account_id: AccountId,
    
    // Staking and delegation information. 
    // `Some` means staking information is available, staking pool selected. 
    // `None` means no staking pool selected.
    pub staking_information: Option<StakingInformation>,

    // AccountId of NEAR Foundation, can terminate vesting. 
    pub foundation_account_id: Option<AccountId>,  
}


impl Default for LockupContract {
    fn default() -> Self {
      env::panic_str("The contract is not initialized.");
    }
}

#[near_bindgen]
impl LockupContract {
    /// Requires 25 TGas
    /// 
    /// Initializes the contract.
    /// (args will be skipped here, explained in types.rs.)
    #[init]
    pub fn new(
      owner_account_id: AccountId,
      lockup_duration: WrappedDuration,
      lockup_timestamp: Option<WrappedTimestamp>,
      transfers_information: TransfersInformation,
      vesting_schedule: Option<VestingScheduleOrHash>,
      release_duration: Option<WrappedDuration>,
      staking_pool_whitelist_account_id: AccountId,
      foundation_account_id: Option<AccountId>,
    ) -> Self {
      require!(
        env::is_valid_account_id(owner_account_id.as_bytes()),
        "Invalid owner's account ID."
      );

      require!(
        env::is_valid_account_id(staking_pool_whitelist_account_id.as_bytes()),
        "Invalid staking pool whitelist's account ID."
      );

      if let TransfersInformation::TransfersDisabled {
        transfer_poll_account_id,
      } = &transfers_information {
        require!(
          env::is_valid_account_id(transfer_poll_account_id.as_bytes()),
          "Invalid transfer poll's account ID."
        );
      };

      let lockup_information = LockupInformation {
        lockup_amount: env::account_balance(),
        termination_withdrawn_tokens: 0,
        lockup_duration: lockup_duration.0,
        release_duration: release_duration.map(|d| d.0),
        lockup_timestamp: lockup_timestamp.map(|d| d.0),
        transfers_information,
      };

      let vesting_information = match vesting_schedule {
        None => {
          require!(
            foundation_account_id.is_none(),
            "Foundation account can't be added without vesting schedule."
          );
          VestingInformation::None
        }

        Some(VestingScheduleOrHash::VestingHash(hash)) => {
          VestingInformation::VestingHash(hash)
        },

        Some(VestingScheduleOrHash::VestingSchedule(vs)) => {
          VestingInformation::VestingSchedule(vs)
        }
      };

      require!(
        vesting_information == VestingInformation::None 
          || env::is_valid_account_id(
            foundation_account_id.as_ref().unwrap().as_bytes()
          ),
        concat!(
          "Either no vesting created or ",
          "Foundation account should be added for vesting schedule."
        )
      );

      Self {
        owner_account_id,
        lockup_information,
        vesting_information,
        staking_information: None,
        staking_pool_whitelist_account_id,
        foundation_account_id,
      }
    }
}

#[cfg(all(test, not(target_arch="wasm32")))]
mod tests {
    use super::*;
    // use std::convert::TryInto;

    use near_sdk::{testing_env, PromiseResult, VMContext};
    // use near_sdk::json_types::U128;

    mod test_utils;
    use test_utils::*;

    // pub type AccountId = String;
    const SALT: [u8; 3] = [1, 2, 3];
    const FAKE_SALT: [u8; 4] = [3, 1, 2, 4];
    const VESTING_CONST: u64 = 10;

    fn basic_context() -> VMContext {
        get_context(
          system_account(), 
          to_yocto(LOCKUP_NEAR), 
          0, 
          to_ts(GENESIS_TIME_IN_DAYS), 
          false,
          // None,
        )
    }


    fn new_vesting_schedule(
      offset_in_days: u64
    ) -> VestingSchedule {
        VestingSchedule {
          start_timestamp: to_ts(GENESIS_TIME_IN_DAYS - YEAR + offset_in_days).into(),
          cliff_timestamp: to_ts(GENESIS_TIME_IN_DAYS + offset_in_days).into(),
          end_timestamp: to_ts(GENESIS_TIME_IN_DAYS + YEAR * 3 + offset_in_days).into(),
        }
    }


    #[allow(dead_code)]
    fn no_vesting_schedule() -> VestingSchedule {
        VestingSchedule {
          start_timestamp: to_ts(0).into(),
          cliff_timestamp: to_ts(0).into(),
          end_timestamp: to_ts(0).into(),
        }
    }


    fn new_contract_with_lockup_duration(
      transfers_enabled: bool,
      vesting_schedule: Option<VestingSchedule>,
      release_duration: Option<WrappedDuration>,
      foundation_account: bool,
      lockup_duration: Duration,
    ) -> LockupContract {
        let lockup_start_information = if transfers_enabled {
          TransfersInformation::TransfersEnabled {
            transfers_timestamp: to_ts(GENESIS_TIME_IN_DAYS).into(),
          }
        } else {
          TransfersInformation::TransfersDisabled {
            transfer_poll_account_id: "transfers".parse().unwrap(),
          }
        };

        let foundation_account_id = if foundation_account {
          Some(account_foundation())
        } else {
          None
        };

        let vesting_schedule = vesting_schedule.map(|vesting_schedule| {
          VestingScheduleOrHash::VestingHash(
            VestingScheduleWithSalt {
              vesting_schedule,
              salt: SALT.to_vec().into(),
            }
            .hash()
            .into(),
          )
        });

        LockupContract::new(
          account_owner(),
          lockup_duration.into(),
          None, 
          lockup_start_information,
          vesting_schedule,
          release_duration,
          "whitelist".parse().unwrap(),
          foundation_account_id,
        )
    }


    fn new_contract(
      transfers_enabled: bool,
      vesting_schedule: Option<VestingSchedule>,
      release_duration: Option<WrappedDuration>,
      foundation_account: bool,
    ) -> LockupContract {
        new_contract_with_lockup_duration(
          transfers_enabled, 
          vesting_schedule, 
          release_duration, 
          foundation_account, 
          to_nanos(YEAR),
        )
    }


    #[allow(dead_code)]
    fn lockup_only_setup() -> (VMContext, LockupContract) {
        let context = basic_context();
        testing_env!(context.clone());
        let contract = new_contract(true, None, None, false);
        (context, contract)
    }

    fn initialize_context() -> VMContext {
      let context = basic_context();
      testing_env!(context.clone());
      context
    }

    fn factory_initialize_context_contract() -> (VMContext, LockupContract) {
      let context = initialize_context();
      let vesting_schedule = new_vesting_schedule(VESTING_CONST);
      let contract = new_contract(true, Some(vesting_schedule), None, true);
      (context, contract)
    }

    
    #[test]
    #[should_panic(expected = "Can only be called by NEAR Foundation")]
    fn test_terminate_vesting_fully_vested() {
      let (mut context, mut contract) = factory_initialize_context_contract();
      context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);

      
      // We changed this. 
      // context.predecessor_account_id = account_foundation().to_string().to_string();
      // to this: 
      context.predecessor_account_id = non_owner().to_string().to_string();

      context.signer_account_id = non_owner().to_string().to_string();
      testing_env!(context.clone());

      let real_vesting = new_vesting_schedule(VESTING_CONST);
      contract.terminate_vesting(Some(VestingScheduleWithSalt {
        vesting_schedule: real_vesting,
        salt: SALT.to_vec().into(),
      }));
    }


    #[test]
    #[should_panic(expected = "Presented vesting schedule and salt don't match the hash.")]
    fn test_different_salt() {
      let mut context = initialize_context();
      let vesting_duration = 10;
      let vesting_schedule = new_vesting_schedule(vesting_duration);
      let mut contract = new_contract(true, Some(vesting_schedule), None, true);
      context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
      context.predecessor_account_id = account_foundation().to_string().to_string();
      context.signer_account_id = non_owner().to_string().to_string();
      testing_env!(context.clone());

      let real_vesting_schedule = new_vesting_schedule(vesting_duration);
      contract.terminate_vesting(Some(VestingScheduleWithSalt {
        vesting_schedule: real_vesting_schedule,
        salt: FAKE_SALT.to_vec().into(),
      }))
    }

    #[test]
    #[should_panic(expected = "Presented vesting schedule and salt don't match the hash.")]
    fn test_different_vesting() {
      let (mut context, mut contract) = factory_initialize_context_contract();
      context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
      context.predecessor_account_id = account_foundation().to_string().to_string();
      context.signer_account_id = non_owner().to_string().to_string();
      testing_env!(context.clone());

      let fake_vesting_schedule = new_vesting_schedule(25);
      contract.terminate_vesting(Some(VestingScheduleWithSalt {
        vesting_schedule: fake_vesting_schedule,
        salt: SALT.to_vec().into(),
      }))
    }

    #[test]
    #[should_panic(expected = "Staking pool is not selected.")]
    fn test_termination_with_staking_without_staking_pool() {
      let lockup_amount = to_yocto(1000);
      let mut context = initialize_context();
      let vesting_schedule = new_vesting_schedule(0);
      let mut contract = new_contract(true, Some(vesting_schedule.clone()), None, true);

      context.is_view = true;
      testing_env!(context.clone());

      // Originally commented out
      // assert_eq!(contract.get_owners_balance().0, to_yocto(0));
      // assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(0));
      // Originally commented out END here

      assert_eq!(contract.get_locked_amount().0, lockup_amount);
      assert_eq!(
        contract.get_unvested_amount(vesting_schedule.clone()).0,
        to_yocto(750)
      );
      assert_eq!(
        contract.get_locked_vested_amount(vesting_schedule.clone()).0,
        to_yocto(250)
      );
      context.is_view = false;

      context.predecessor_account_id = account_owner().to_string().to_string();
      context.signer_account_pk = public_key(1).into_bytes();
      testing_env!(context.clone());

      // Selecting staking pool
      // --skipped--

      context.is_view = false;
      context.predecessor_account_id = account_foundation().to_string().to_string();
      context.signer_account_pk = public_key(2).into_bytes();
      testing_env!(context.clone());
      contract.termination_prepare_to_withdraw();
      assert_eq!(
        contract.get_termination_status(),
        Some(TerminationStatus::UnstakingInProgress)
      );

    }

    // ============================= OTHER TESTS ============================ //
    #[test]
    fn test_lockup_only_basic() {
        let (mut context, contract) = lockup_only_setup();
        // Checking initial values at genesis time
        context.is_view = true;
        testing_env!(context.clone());

        assert_eq!(contract.get_owners_balance().0, 0);
        assert_eq!(
            contract.get_locked_vested_amount(no_vesting_schedule()).0,
            to_yocto(LOCKUP_NEAR)
        );

        // Checking values in 1 day after genesis time
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 1);

        assert_eq!(contract.get_owners_balance().0, 0);

        // Checking values next day after lockup timestamp
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR + 1);
        testing_env!(context.clone());

        assert_almost_eq(contract.get_owners_balance().0, to_yocto(LOCKUP_NEAR));
    }

    #[test]
    fn test_add_full_access_key() {
        let (mut context, mut contract) = lockup_only_setup();
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(1).try_into().unwrap();
        testing_env!(context.clone());

        contract.add_full_access_key(public_key(4));
    }

    #[test]
    #[should_panic(expected = "Tokens are still locked/unvested")]
    fn test_add_full_access_key_when_vesting_is_not_finished() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let vesting_schedule = new_vesting_schedule(YEAR);
        let mut contract = new_contract(true, Some(vesting_schedule), None, true);

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR - 10);
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(1).try_into().unwrap();
        testing_env!(context.clone());

        contract.add_full_access_key(public_key(4));
    }

    #[test]
    #[should_panic(expected = "Tokens are still locked/unvested")]
    fn test_add_full_access_key_when_lockup_is_not_finished() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let mut contract = new_contract(true, None, Some(to_nanos(YEAR).into()), false);

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR - 10);
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(1).try_into().unwrap();
        testing_env!(context.clone());

        contract.add_full_access_key(public_key(4));
    }

    #[test]
    #[should_panic(expected = "This method can only be called by the owner. ")]
    fn test_call_by_non_owner() {
        let (mut context, mut contract) = lockup_only_setup();
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
        context.predecessor_account_id = non_owner().to_string();
        context.signer_account_id = non_owner().to_string();
        testing_env!(context.clone());

        contract.select_staking_pool("staking_pool".parse().unwrap());
    }

    #[test]
    #[should_panic(expected = "Presented vesting schedule and salt don't match the hash")]
    fn test_vesting_doesnt_match() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let vesting_schedule = new_vesting_schedule(5);
        let mut contract = new_contract(true, Some(vesting_schedule), None, true);
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
        context.predecessor_account_id = account_foundation().to_string();
        context.signer_account_id = non_owner().to_string();
        testing_env!(context.clone());

        let not_real_vesting = new_vesting_schedule(100);
        contract.terminate_vesting(Some(VestingScheduleWithSalt {
            vesting_schedule: not_real_vesting,
            salt: SALT.to_vec().into(),
        }));
    }

    #[test]
    #[should_panic(expected = "called `Result::unwrap()` on an `Err` value: HostError(GuestPanic { panic_msg: \"Expected vesting schedule and salt, but not provided.\" })")]
    fn test_vesting_schedule_and_salt_not_provided() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let vesting_schedule = new_vesting_schedule(5);
        let mut contract = new_contract(true, Some(vesting_schedule), None, true);
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
        context.predecessor_account_id = account_foundation().to_string();
        context.signer_account_id = non_owner().to_string();
        testing_env!(context.clone());

        contract.terminate_vesting(None);
    }

    #[test]
    #[should_panic(expected = "Explicit vesting schedule already exists.")]
    fn test_explicit_vesting() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let vesting_schedule = new_vesting_schedule(5);
        let mut contract = LockupContract::new(
            account_owner(),
            to_nanos(YEAR).into(),
            None,
            TransfersInformation::TransfersEnabled {
                transfers_timestamp: to_ts(GENESIS_TIME_IN_DAYS).into(),
            },
            Some(VestingScheduleOrHash::VestingSchedule(
                vesting_schedule.clone(),
            )),
            None,
            "whitelist".parse().unwrap(),
            Some(account_foundation()),
        );
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
        context.predecessor_account_id = account_foundation().to_string();
        context.signer_account_id = non_owner().to_string();
        testing_env!(context.clone());

        contract.terminate_vesting(Some(VestingScheduleWithSalt {
            vesting_schedule,
            salt: SALT.to_vec().into(),
        }));
    }

    #[test]
    #[should_panic(expected = "Foundation account can't be added without vesting schedule")]
    fn test_init_foundation_key_no_vesting() {
        let context = basic_context();
        testing_env!(context.clone());
        new_contract(true, None, None, true);
    }

    #[test]
    #[should_panic(expected = "Foundation account can't be added without vesting schedule")]
    fn test_init_foundation_key_no_vesting_with_release() {
        let context = basic_context();
        testing_env!(context.clone());
        new_contract(true, None, Some(to_nanos(YEAR).into()), true);
    }

    #[test]
    #[should_panic(expected = "Can only be called by NEAR Foundation")]
    fn test_call_by_non_foundation() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let vesting_schedule = new_vesting_schedule(0);
        let mut contract = new_contract(true, Some(vesting_schedule.clone()), None, true);
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
        context.predecessor_account_id = non_owner().to_string();
        context.signer_account_id = non_owner().to_string();
        testing_env!(context.clone());

        contract.terminate_vesting(None);
    }

    #[test]
    #[should_panic(expected = "Transfers are disabled")]
    fn test_transfers_not_enabled() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let mut contract = new_contract(false, None, None, false);
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR + 1);
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(1).try_into().unwrap();
        context.is_view = false;
        testing_env!(context.clone());

        contract.transfer(to_yocto(100).into(), non_owner());
    }

    #[test]
    fn test_enable_transfers() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let mut contract = new_contract(false, None, None, false);
        context.is_view = true;
        testing_env!(context.clone());
        assert!(!contract.are_transfers_enabled());

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR + 1);
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(1).try_into().unwrap();
        context.is_view = false;
        testing_env!(context.clone());

        contract.check_transfers_vote();

        let poll_result = Some(to_ts(GENESIS_TIME_IN_DAYS + 10).into());
        context.predecessor_account_id = lockup_account().to_string();
        // NOTE: Unit tests don't need to read the content of the promise result. So here we don't
        // have to pass serialized result from the transfer poll.
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        assert!(contract.on_get_result_from_transfer_poll(poll_result));

        context.is_view = true;
        testing_env!(context.clone());
        // Not unlocked yet
        assert_eq!(contract.get_owners_balance().0, 0);
        assert!(contract.are_transfers_enabled());
        assert_eq!(contract.get_vesting_information(), VestingInformation::None);

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR + 10);
        testing_env!(context.clone());
        // Unlocked yet
        assert_eq!(
            contract.get_owners_balance().0,
            to_yocto(LOCKUP_NEAR).into()
        );

        context.is_view = false;
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.transfer(to_yocto(100).into(), non_owner());
    }

    #[test]
    fn test_check_transfers_vote_false() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let mut contract = new_contract(false, None, None, false);
        context.is_view = true;
        testing_env!(context.clone());
        assert!(!contract.are_transfers_enabled());

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR + 1);
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(1).try_into().unwrap();
        context.is_view = false;
        testing_env!(context.clone());

        contract.check_transfers_vote();

        let poll_result = None;
        // NOTE: Unit tests don't need to read the content of the promise result. So here we don't
        // have to pass serialized result from the transfer poll.
        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        assert!(!contract.on_get_result_from_transfer_poll(poll_result));

        context.is_view = true;
        testing_env!(context.clone());
        // Not enabled
        assert!(!contract.are_transfers_enabled());
    }

    #[test]
    fn test_lockup_only_transfer_call_by_owner() {
        let (mut context, mut contract) = lockup_only_setup();
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR + 1);
        context.is_view = true;
        testing_env!(context.clone());
        assert_almost_eq(contract.get_owners_balance().0, to_yocto(LOCKUP_NEAR));

        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(1).try_into().unwrap();
        context.is_view = false;
        testing_env!(context.clone());

        assert_eq!(env::account_balance(), to_yocto(LOCKUP_NEAR));
        contract.transfer(to_yocto(100).into(), non_owner());
        assert_almost_eq(env::account_balance(), to_yocto(LOCKUP_NEAR - 100));
    }

    #[test]
    #[should_panic(expected = "Staking pool is not selected")]
    fn test_staking_pool_is_not_selected() {
        let (mut context, mut contract) = lockup_only_setup();
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(2).try_into().unwrap();

        let amount = to_yocto(LOCKUP_NEAR - 100);
        testing_env!(context.clone());
        contract.deposit_to_staking_pool(amount.into());
    }

    #[test]
    fn test_staking_pool_success() {
        let (mut context, mut contract) = lockup_only_setup();
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(2).try_into().unwrap();

        // Selecting staking pool
        let staking_pool: AccountId = "staking_pool".parse().unwrap();
        testing_env!(context.clone());
        contract.select_staking_pool(staking_pool.clone());

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(
            context.clone(),
            PromiseResult::Successful(b"true".to_vec()),
        );
        contract.on_whitelist_is_whitelisted(true, staking_pool.clone());

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_staking_pool_account_id(), Some(staking_pool));
        assert_eq!(contract.get_known_deposited_balance().0, 0);
        context.is_view = false;

        // Deposit to the staking_pool
        let amount = to_yocto(LOCKUP_NEAR - 100);
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.deposit_to_staking_pool(amount.into());
        context.account_balance = env::account_balance();
        assert_eq!(context.account_balance, to_yocto(LOCKUP_NEAR) - amount);

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_staking_pool_deposit(amount.into());
        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_known_deposited_balance().0, amount);
        context.is_view = false;

        // Staking on the staking pool
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.stake(amount.into());

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_staking_pool_stake(amount.into());

        // Assuming there are 20 NEAR tokens in rewards. Unstaking.
        let unstake_amount = amount + to_yocto(20);
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.unstake(unstake_amount.into());

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_staking_pool_unstake(unstake_amount.into());

        // Withdrawing
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.withdraw_from_staking_pool(unstake_amount.into());
        context.account_balance += unstake_amount;

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_staking_pool_withdraw(unstake_amount.into());
        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_known_deposited_balance().0, 0);
        context.is_view = false;

        // Unselecting staking pool
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.unselect_staking_pool();
        assert_eq!(contract.get_staking_pool_account_id(), None);
    }

    #[test]
    fn test_staking_pool_refresh_balance() {
        let (mut context, mut contract) = lockup_only_setup();
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(2).try_into().unwrap();

        // Selecting staking pool
        let staking_pool: AccountId = "staking_pool".parse().unwrap();
        testing_env!(context.clone());
        contract.select_staking_pool(staking_pool.clone());

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(
            context.clone(),
            PromiseResult::Successful(b"true".to_vec()),
        );
        contract.on_whitelist_is_whitelisted(true, staking_pool.clone());

        // Deposit to the staking_pool
        let amount = to_yocto(LOCKUP_NEAR - 100);
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.deposit_to_staking_pool(amount.into());
        context.account_balance = env::account_balance();
        assert_eq!(context.account_balance, to_yocto(LOCKUP_NEAR) - amount);

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_staking_pool_deposit(amount.into());

        // Staking on the staking pool
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.stake(amount.into());

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_staking_pool_stake(amount.into());

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, 0);
        assert_eq!(contract.get_liquid_owners_balance().0, 0);
        assert_eq!(contract.get_known_deposited_balance().0, amount);
        context.is_view = false;

        // Assuming there are 20 NEAR tokens in rewards. Refreshing balance.
        let total_balance = amount + to_yocto(20);
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.refresh_staking_pool_balance();

        // In unit tests, the following call ignores the promise value, because it's passed directly.
        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_get_account_total_balance(total_balance.into());

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_known_deposited_balance().0, total_balance);
        assert_eq!(contract.get_owners_balance().0, to_yocto(20));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(20));
        context.is_view = false;

        // Withdrawing these tokens
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        let transfer_amount = to_yocto(15);
        contract.transfer(transfer_amount.into(), non_owner());
        context.account_balance = env::account_balance();

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_known_deposited_balance().0, total_balance);
        assert_eq!(contract.get_owners_balance().0, to_yocto(5));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(5));
        context.is_view = false;
    }

    // ================================= PART 2 ===================================== //
    #[test]
    #[should_panic(expected = "Staking pool is already selected")]
    fn test_staking_pool_selected_again() {
        let (mut context, mut contract) = lockup_only_setup();
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(2).try_into().unwrap();

        // Selecting staking pool
        let staking_pool: AccountId = "staking_pool".parse().unwrap();
        testing_env!(context.clone());
        contract.select_staking_pool(staking_pool.clone());

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(
            context.clone(),
            PromiseResult::Successful(b"true".to_vec()),
        );
        contract.on_whitelist_is_whitelisted(true, staking_pool.clone());

        // Selecting another staking pool
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.select_staking_pool("staking_pool_2".parse().unwrap());
    }

    #[test]
    #[should_panic(expected = "The given staking pool ID is not whitelisted.")]
    fn test_staking_pool_not_whitelisted() {
        let (mut context, mut contract) = lockup_only_setup();
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(2).try_into().unwrap();

        // Selecting staking pool
        let staking_pool: AccountId = "staking_pool".parse().unwrap();
        testing_env!(context.clone());
        contract.select_staking_pool(staking_pool.clone());

        context.predecessor_account_id = lockup_account().to_string();
        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(
            context.clone(),
            PromiseResult::Successful(b"false".to_vec()),
        );
        contract.on_whitelist_is_whitelisted(false, staking_pool.clone());
    }

    #[test]
    #[should_panic(expected = "Staking pool is not selected")]
    fn test_staking_pool_unselecting_non_selected() {
        let (mut context, mut contract) = lockup_only_setup();
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(2).try_into().unwrap();

        // Unselecting staking pool
        testing_env!(context.clone());
        contract.unselect_staking_pool();
    }

    #[test]
    #[should_panic(expected = "There is still deposit on staking pool.")]
    fn test_staking_pool_unselecting_with_deposit() {
        let (mut context, mut contract) = lockup_only_setup();
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(2).try_into().unwrap();

        // Selecting staking pool
        let staking_pool: AccountId = "staking_pool".parse().unwrap();
        testing_env!(context.clone());
        contract.select_staking_pool(staking_pool.clone());

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(
            context.clone(),
            PromiseResult::Successful(b"true".to_vec()),
        );
        contract.on_whitelist_is_whitelisted(true, staking_pool.clone());

        // Deposit to the staking_pool
        let amount = to_yocto(LOCKUP_NEAR - 100);
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.deposit_to_staking_pool(amount.into());
        context.account_balance = env::account_balance();

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_staking_pool_deposit(amount.into());

        // Unselecting staking pool
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.unselect_staking_pool();
    }

    #[test]
    fn test_staking_pool_owner_balance() {
        let (mut context, mut contract) = lockup_only_setup();
        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(2).try_into().unwrap();
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR + 1);

        let lockup_amount = to_yocto(LOCKUP_NEAR);
        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, lockup_amount);
        context.is_view = false;

        // Selecting staking pool
        let staking_pool: AccountId = "staking_pool".parse().unwrap();
        testing_env!(context.clone());
        contract.select_staking_pool(staking_pool.clone());

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(
            context.clone(),
            PromiseResult::Successful(b"true".to_vec()),
        );
        contract.on_whitelist_is_whitelisted(true, staking_pool.clone());

        // Deposit to the staking_pool
        let mut total_amount = 0;
        let amount = to_yocto(100);
        for _ in 1..=5 {
            total_amount += amount;
            context.predecessor_account_id = account_owner().to_string();
            testing_env!(context.clone());
            contract.deposit_to_staking_pool(amount.into());
            context.account_balance = env::account_balance();
            assert_eq!(context.account_balance, lockup_amount - total_amount);

            context.predecessor_account_id = lockup_account().to_string();
            testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
            contract.on_staking_pool_deposit(amount.into());
            context.is_view = true;
            testing_env!(context.clone());
            assert_eq!(contract.get_known_deposited_balance().0, total_amount);
            assert_eq!(contract.get_owners_balance().0, lockup_amount);
            assert_eq!(
                contract.get_liquid_owners_balance().0,
                lockup_amount - total_amount - MIN_BALANCE_FOR_STORAGE
            );
            context.is_view = false;
        }

        // Withdrawing from the staking_pool. Plus one extra time as a reward
        let mut total_withdrawn_amount = 0;
        for _ in 1..=6 {
            total_withdrawn_amount += amount;
            context.predecessor_account_id = account_owner().to_string();
            testing_env!(context.clone());
            contract.withdraw_from_staking_pool(amount.into());
            context.account_balance += amount;
            assert_eq!(
                context.account_balance,
                lockup_amount - total_amount + total_withdrawn_amount
            );

            context.predecessor_account_id = lockup_account().to_string();
            testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
            contract.on_staking_pool_withdraw(amount.into());
            context.is_view = true;
            testing_env!(context.clone());
            assert_eq!(
                contract.get_known_deposited_balance().0,
                total_amount.saturating_sub(total_withdrawn_amount)
            );
            assert_eq!(
                contract.get_owners_balance().0,
                lockup_amount + total_withdrawn_amount.saturating_sub(total_amount)
            );
            assert_eq!(
                contract.get_liquid_owners_balance().0,
                lockup_amount - total_amount + total_withdrawn_amount - MIN_BALANCE_FOR_STORAGE
            );
            context.is_view = false;
        }
    }

    #[test]
    fn test_lock_timestmap() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let contract = LockupContract::new(
            account_owner(),
            0.into(),
            Some(to_ts(GENESIS_TIME_IN_DAYS + YEAR).into()),
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfers".parse().unwrap(),
            },
            None,
            None,
            "whitelist".parse().unwrap(),
            None,
        );

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, 0);
        assert_eq!(contract.get_liquid_owners_balance().0, 0);
        assert_eq!(
            contract.get_locked_vested_amount(no_vesting_schedule()).0,
            to_yocto(1000)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(1000));
        assert_eq!(
            contract.get_unvested_amount(no_vesting_schedule()).0,
            to_yocto(0)
        );
        assert!(!contract.are_transfers_enabled());

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, 0);
        assert_eq!(contract.get_liquid_owners_balance().0, 0);
        assert_eq!(
            contract.get_locked_vested_amount(no_vesting_schedule()).0,
            to_yocto(1000)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(1000));
        assert_eq!(
            contract.get_unvested_amount(no_vesting_schedule()).0,
            to_yocto(0)
        );
    }

    #[test]
    fn test_lock_timestmap_transfer_enabled() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let contract = LockupContract::new(
            account_owner(),
            0.into(),
            Some(to_ts(GENESIS_TIME_IN_DAYS + YEAR).into()),
            TransfersInformation::TransfersEnabled {
                transfers_timestamp: to_ts(GENESIS_TIME_IN_DAYS + YEAR / 2).into(),
            },
            None,
            None,
            "whitelist".parse().unwrap(),
            None,
        );

        context.is_view = true;
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(1000));
        assert_eq!(
            contract.get_liquid_owners_balance().0,
            to_yocto(1000) - MIN_BALANCE_FOR_STORAGE
        );
        assert_eq!(
            contract.get_locked_vested_amount(no_vesting_schedule()).0,
            to_yocto(0)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(0));
        assert_eq!(
            contract.get_unvested_amount(no_vesting_schedule()).0,
            to_yocto(0)
        );
    }

    #[test]
    fn test_termination_no_staking() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let vesting_schedule = new_vesting_schedule(0);
        let mut contract = LockupContract::new(
            account_owner(),
            to_nanos(YEAR).into(),
            None,
            TransfersInformation::TransfersEnabled {
                transfers_timestamp: to_ts(GENESIS_TIME_IN_DAYS).into(),
            },
            Some(VestingScheduleOrHash::VestingSchedule(
                vesting_schedule.clone(),
            )),
            None,
            "whitelist".parse().unwrap(),
            Some(account_foundation()),
        );

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(
            contract.get_vesting_information(),
            VestingInformation::VestingSchedule(vesting_schedule.clone())
        );
        assert_eq!(contract.get_owners_balance().0, 0);
        assert_eq!(contract.get_liquid_owners_balance().0, 0);
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(1000));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(750)
        );

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(500));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(500));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(0)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(500));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(500)
        );

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 2 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(750));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(750));
        assert_eq!(contract.get_locked_amount().0, to_yocto(250));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(250)
        );

        // Terminating
        context.is_view = false;
        context.predecessor_account_id = account_foundation().to_string();
        context.signer_account_pk = public_key(3).into();
        testing_env!(context.clone());
        contract.terminate_vesting(None);

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(
            contract.get_vesting_information(),
            VestingInformation::Terminating(TerminationInformation {
                unvested_amount: to_yocto(250).into(),
                status: TerminationStatus::ReadyToWithdraw,
            })
        );
        assert_eq!(contract.get_owners_balance().0, to_yocto(750));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(750));
        assert_eq!(contract.get_locked_amount().0, to_yocto(250));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(250)
        );
        assert_eq!(contract.get_terminated_unvested_balance().0, to_yocto(250));
        assert_eq!(
            contract.get_terminated_unvested_balance_deficit().0,
            to_yocto(0)
        );
        assert_eq!(
            contract.get_termination_status(),
            Some(TerminationStatus::ReadyToWithdraw)
        );

        // Withdrawing
        context.is_view = false;
        testing_env!(context.clone());
        let receiver_id: AccountId = "near".parse().unwrap();
        contract.termination_withdraw(receiver_id.clone());
        context.account_balance = env::account_balance();

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_withdraw_unvested_amount(to_yocto(250).into(), receiver_id);

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(750));
        assert_eq!(
            contract.get_liquid_owners_balance().0,
            to_yocto(750) - MIN_BALANCE_FOR_STORAGE
        );
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(0)
        );
        assert_eq!(contract.get_terminated_unvested_balance().0, to_yocto(0));
        assert_eq!(contract.get_termination_status(), None);
        assert_eq!(contract.get_vesting_information(), VestingInformation::None);
    }

    #[test]
    fn test_release_duration() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let contract = new_contract(true, None, Some(to_nanos(4 * YEAR).into()), false);

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, 0);
        assert_eq!(contract.get_liquid_owners_balance().0, 0);
        assert_eq!(
            contract.get_locked_vested_amount(no_vesting_schedule()).0,
            to_yocto(1000)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(1000));
        assert_eq!(
            contract.get_unvested_amount(no_vesting_schedule()).0,
            to_yocto(0)
        );

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(0));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(0));
        assert_eq!(
            contract.get_locked_vested_amount(no_vesting_schedule()).0,
            to_yocto(1000)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(1000));

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 2 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(250));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(250));
        assert_eq!(
            contract.get_locked_vested_amount(no_vesting_schedule()).0,
            to_yocto(750)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(750));

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 3 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(500));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(500));
        assert_eq!(
            contract.get_locked_vested_amount(no_vesting_schedule()).0,
            to_yocto(500)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(500));

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 4 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(750));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(750));
        assert_eq!(contract.get_locked_amount().0, to_yocto(250));
    }


    // ================================= PART 3 ==================================== //
    #[test]
    fn test_vesting_and_release_duration() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let vesting_schedule = new_vesting_schedule(0);
        let contract = new_contract_with_lockup_duration(
            true,
            Some(vesting_schedule.clone()),
            Some(to_nanos(4 * YEAR).into()),
            true,
            0,
        );

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, 0);
        assert_eq!(contract.get_liquid_owners_balance().0, 0);
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(1000));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(750)
        );

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(250));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(250));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(750));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(500)
        );

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 2 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(500));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(500));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(500));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(250)
        );

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 3 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(750));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(750));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(250));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(0)
        );

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 4 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(1000));
        assert_eq!(
            contract.get_liquid_owners_balance().0,
            to_yocto(1000) - MIN_BALANCE_FOR_STORAGE
        );
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(0)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(0));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(0)
        );
    }

    // Vesting post transfers is not supported by Hash vesting.
    #[test]
    fn test_vesting_post_transfers_and_release_duration() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let vesting_schedule = new_vesting_schedule(YEAR * 2);
        let contract = LockupContract::new(
            account_owner(),
            to_nanos(YEAR).into(),
            None,
            TransfersInformation::TransfersEnabled {
                transfers_timestamp: to_ts(GENESIS_TIME_IN_DAYS).into(),
            },
            Some(VestingScheduleOrHash::VestingSchedule(
                vesting_schedule.clone(),
            )),
            Some(to_nanos(4 * YEAR).into()),
            "whitelist".parse().unwrap(),
            Some(account_foundation()),
        );

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, 0);
        assert_eq!(contract.get_liquid_owners_balance().0, 0);
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(1000)
        );
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(0)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(1000));

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(0));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(0));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(0)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(1000));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(1000)
        );

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 2 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(250));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(250));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(0)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(750));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(750)
        );

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 3 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(500));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(500));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(0)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(500));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(500)
        );

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 4 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(750));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(750));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(0)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(250));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(250)
        );

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 5 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(1000));
        assert_eq!(
            contract.get_liquid_owners_balance().0,
            to_yocto(1000) - MIN_BALANCE_FOR_STORAGE
        );
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(0)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(0));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(0)
        );
    }

    #[test]
    fn test_termination_no_staking_with_release_duration() {
        let mut context = basic_context();
        testing_env!(context.clone());
        let vesting_schedule = new_vesting_schedule(0);
        let mut contract = new_contract_with_lockup_duration(
            true,
            Some(vesting_schedule.clone()),
            Some(to_nanos(4 * YEAR).into()),
            true,
            0,
        );

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, 0);
        assert_eq!(contract.get_liquid_owners_balance().0, 0);
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(1000));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(750)
        );

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 2 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(500));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(500));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(500));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(250)
        );

        // Terminating
        context.is_view = false;
        context.predecessor_account_id = account_foundation().to_string();
        context.signer_account_pk = public_key(3).into();
        testing_env!(context.clone());
        contract.terminate_vesting(Some(VestingScheduleWithSalt {
            vesting_schedule: vesting_schedule.clone(),
            salt: SALT.to_vec().into(),
        }));

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(500));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(500));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(contract.get_locked_amount().0, to_yocto(500));
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(250)
        );
        assert_eq!(contract.get_terminated_unvested_balance().0, to_yocto(250));
        assert_eq!(
            contract.get_terminated_unvested_balance_deficit().0,
            to_yocto(0)
        );
        assert_eq!(
            contract.get_termination_status(),
            Some(TerminationStatus::ReadyToWithdraw)
        );

        // Withdrawing
        context.is_view = false;
        testing_env!(context.clone());
        let receiver_id: AccountId = "near".parse().unwrap();
        contract.termination_withdraw(receiver_id.clone());
        context.account_balance = env::account_balance();

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_withdraw_unvested_amount(to_yocto(250).into(), receiver_id);

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(500));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(500));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(0)
        );
        assert_eq!(contract.get_terminated_unvested_balance().0, to_yocto(0));
        assert_eq!(contract.get_termination_status(), None);

        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + 3 * YEAR);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(750));
        assert_eq!(
            contract.get_liquid_owners_balance().0,
            to_yocto(750) - MIN_BALANCE_FOR_STORAGE
        );
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(0)
        );
    }

    #[test]
    fn test_termination_before_cliff() {
        let lockup_amount = to_yocto(1000);
        let mut context = basic_context();
        testing_env!(context.clone());
        let vesting_schedule = new_vesting_schedule(YEAR);
        let mut contract = new_contract(true, Some(vesting_schedule.clone()), None, true);

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(
            contract.get_vesting_information(),
            VestingInformation::VestingHash(
                VestingScheduleWithSalt {
                    vesting_schedule: vesting_schedule.clone(),
                    salt: SALT.to_vec().into(),
                }
                    .hash()
                    .into()
            )
        );
        assert_eq!(contract.get_owners_balance().0, 0);
        assert_eq!(contract.get_liquid_owners_balance().0, 0);
        assert_eq!(contract.get_locked_amount().0, lockup_amount);
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            lockup_amount
        );
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            0
        );

        // Terminating
        context.is_view = false;
        context.predecessor_account_id = account_foundation().to_string();
        context.signer_account_pk = public_key(3).into();
        testing_env!(context.clone());
        contract.terminate_vesting(Some(VestingScheduleWithSalt {
            vesting_schedule: vesting_schedule.clone(),
            salt: SALT.to_vec().into(),
        }));

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(
            contract.get_vesting_information(),
            VestingInformation::Terminating(TerminationInformation {
                unvested_amount: lockup_amount.into(),
                status: TerminationStatus::ReadyToWithdraw,
            })
        );
        assert_eq!(contract.get_owners_balance().0, 0);
        assert_eq!(contract.get_liquid_owners_balance().0, 0);
        assert_eq!(contract.get_locked_amount().0, lockup_amount);
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            lockup_amount
        );
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            0
        );
        assert_eq!(contract.get_terminated_unvested_balance().0, lockup_amount);
        assert_eq!(
            contract.get_terminated_unvested_balance_deficit().0,
            MIN_BALANCE_FOR_STORAGE
        );
        assert_eq!(
            contract.get_termination_status(),
            Some(TerminationStatus::ReadyToWithdraw)
        );

        // Withdrawing
        context.is_view = false;
        testing_env!(context.clone());
        let receiver_id = account_foundation();
        contract.termination_withdraw(receiver_id.clone());
        context.account_balance = env::account_balance();
        assert_eq!(context.account_balance, MIN_BALANCE_FOR_STORAGE);

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_withdraw_unvested_amount(
            (lockup_amount - MIN_BALANCE_FOR_STORAGE).into(),
            receiver_id,
        );

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            MIN_BALANCE_FOR_STORAGE
        );
        assert_eq!(contract.get_owners_balance().0, 0);
        assert_eq!(contract.get_liquid_owners_balance().0, 0);
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            0
        );
        assert_eq!(
            contract.get_terminated_unvested_balance().0,
            MIN_BALANCE_FOR_STORAGE
        );
        assert_eq!(
            contract.get_terminated_unvested_balance_deficit().0,
            MIN_BALANCE_FOR_STORAGE
        );
        assert_eq!(
            contract.get_termination_status(),
            Some(TerminationStatus::ReadyToWithdraw)
        );
    }

    #[test]
    fn test_termination_with_staking() {
        let lockup_amount = to_yocto(1000);
        let mut context = basic_context();
        testing_env!(context.clone());
        let vesting_schedule = new_vesting_schedule(0);
        let mut contract = new_contract(true, Some(vesting_schedule.clone()), None, true);

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(0));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(0));
        assert_eq!(contract.get_locked_amount().0, lockup_amount);
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(750)
        );
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        context.is_view = false;

        context.predecessor_account_id = account_owner().to_string();
        context.signer_account_pk = public_key(2).into();
        testing_env!(context.clone());

        // Selecting staking pool
        let staking_pool: AccountId = "staking_pool".parse().unwrap();
        testing_env!(context.clone());
        contract.select_staking_pool(staking_pool.clone());

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(
            context.clone(),
            PromiseResult::Successful(b"true".to_vec()),
        );
        contract.on_whitelist_is_whitelisted(true, staking_pool.clone());

        // Deposit to the staking_pool
        let stake_amount = to_yocto(LOCKUP_NEAR - 100);
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.deposit_to_staking_pool(stake_amount.into());
        context.account_balance = env::account_balance();

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_staking_pool_deposit(stake_amount.into());

        // Staking on the staking pool
        context.predecessor_account_id = account_owner().to_string();
        testing_env!(context.clone());
        contract.stake(stake_amount.into());

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_staking_pool_stake(stake_amount.into());

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(0));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(0));
        assert_eq!(contract.get_known_deposited_balance().0, stake_amount);
        assert_eq!(contract.get_locked_amount().0, lockup_amount);
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(750)
        );
        context.is_view = false;

        // Foundation terminating
        context.is_view = false;
        context.predecessor_account_id = account_foundation().to_string();
        context.signer_account_pk = public_key(3).into();
        testing_env!(context.clone());
        contract.terminate_vesting(Some(VestingScheduleWithSalt {
            vesting_schedule: vesting_schedule.clone(),
            salt: SALT.to_vec().into(),
        }));

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(0));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(0));
        assert_eq!(contract.get_locked_amount().0, lockup_amount);
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(750)
        );
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(contract.get_terminated_unvested_balance().0, to_yocto(750));
        assert_eq!(
            contract.get_terminated_unvested_balance_deficit().0,
            to_yocto(650) + MIN_BALANCE_FOR_STORAGE
        );
        assert_eq!(
            contract.get_termination_status(),
            Some(TerminationStatus::VestingTerminatedWithDeficit)
        );

        // Proceeding with unstaking from the pool due to termination.
        context.is_view = false;
        testing_env!(context.clone());
        contract.termination_prepare_to_withdraw();
        assert_eq!(
            contract.get_termination_status(),
            Some(TerminationStatus::UnstakingInProgress)
        );

        let stake_amount_with_rewards = stake_amount + to_yocto(50);
        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(
            context.clone(),
            PromiseResult::Successful(format!("{}", stake_amount_with_rewards).into_bytes()),
        );
        contract.on_get_account_staked_balance_to_unstake(stake_amount_with_rewards.into());

        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_staking_pool_unstake_for_termination(stake_amount_with_rewards.into());

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(
            contract.get_termination_status(),
            Some(TerminationStatus::EverythingUnstaked)
        );

        // Proceeding with withdrawing from the pool due to termination.
        context.is_view = false;
        context.predecessor_account_id = account_foundation().to_string();
        testing_env!(context.clone());
        contract.termination_prepare_to_withdraw();
        assert_eq!(
            contract.get_termination_status(),
            Some(TerminationStatus::WithdrawingFromStakingPoolInProgress)
        );

        let withdraw_amount_with_extra_rewards = stake_amount_with_rewards + to_yocto(1);
        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(
            context.clone(),
            PromiseResult::Successful(
                format!("{}", withdraw_amount_with_extra_rewards).into_bytes(),
            ),
        );
        contract
            .on_get_account_unstaked_balance_to_withdraw(withdraw_amount_with_extra_rewards.into());
        context.account_balance += withdraw_amount_with_extra_rewards;

        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract
            .on_staking_pool_withdraw_for_termination(withdraw_amount_with_extra_rewards.into());

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(51));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(51));
        assert_eq!(contract.get_locked_amount().0, lockup_amount);
        assert_eq!(
            contract.get_unvested_amount(vesting_schedule.clone()).0,
            to_yocto(750)
        );
        assert_eq!(contract.get_terminated_unvested_balance().0, to_yocto(750));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(contract.get_terminated_unvested_balance_deficit().0, 0);
        assert_eq!(contract.get_known_deposited_balance().0, 0);
        assert_eq!(
            contract.get_termination_status(),
            Some(TerminationStatus::ReadyToWithdraw)
        );

        // Withdrawing
        context.is_view = false;
        context.predecessor_account_id = account_foundation().to_string();
        testing_env!(context.clone());
        let receiver_id = account_foundation();
        contract.termination_withdraw(receiver_id.clone());
        context.account_balance = env::account_balance();
        assert_eq!(context.account_balance, to_yocto(250 + 51));

        context.predecessor_account_id = lockup_account().to_string();
        testing_env_with_promise_results(context.clone(), PromiseResult::Successful(vec![]));
        contract.on_withdraw_unvested_amount(to_yocto(750).into(), receiver_id);

        context.is_view = true;
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(51));
        assert_eq!(contract.get_liquid_owners_balance().0, to_yocto(51));
        assert_eq!(contract.get_locked_amount().0, to_yocto(250));
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            to_yocto(250)
        );
        assert_eq!(contract.get_unvested_amount(vesting_schedule.clone()).0, 0);
        assert_eq!(contract.get_terminated_unvested_balance().0, 0);
        assert_eq!(contract.get_terminated_unvested_balance_deficit().0, 0);
        assert_eq!(contract.get_termination_status(), None);

        // Checking the balance becomes unlocked later
        context.block_timestamp = to_ts(GENESIS_TIME_IN_DAYS + YEAR + 1);
        testing_env!(context.clone());
        assert_eq!(contract.get_owners_balance().0, to_yocto(301));
        assert_eq!(
            contract.get_liquid_owners_balance().0,
            to_yocto(301) - MIN_BALANCE_FOR_STORAGE
        );
        assert_eq!(
            contract
                .get_locked_vested_amount(vesting_schedule.clone())
                .0,
            0
        );
        assert_eq!(contract.get_locked_amount().0, 0);
    }
}

Next, we check that the owner currently does not have anything staked on the staking pool.

use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

The output is also interesting. We have an .view_method_call which just performs something similar to near view call from near-cli. The returned result is a JSON, so we call unwrap_json() to that. Now, it's not necessarily something returned; if there's nothing, then None is returned instead. In reality, the contract will show an error message with SmartContractPanic error.

We move on to selecting the staking pool and check they are selected as we expect them to.

use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

We use a function_call here, because we're trying to imitate near call. In short, the near-cli version of this function is:

near call $LOCKUP_CONTRACT select_staking_pool '{
  "account_id": "$STAKING_POOL_ACCOUNT_ID"
}' --gas=$MAX_GAS --amount=0 --accountId $OWNER_STAKING_ACCOUNT

assuming we store the values as environment variables already. We only select the staking pool, nothing is deposited yet; we now have a staking_pool_account_id and we can view it with near view $LOCKUP_CONTRACT get_staking_pool_account_id '{}'; but we don't have any balance (as seen in the last line, res.0 == 0).

We shall now make deposits to the staking pool. We will deposit 100 NEAR.

use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

Previously we mentioned that staking involves two operation: depositing the fund to the corresponding staking account first, before it's being stake. Here, it's also 2 separate function calls. We have deposited to the staking pool, but not yet stake it. We can see our deposit with get_known_deposited_balance view function. We shall stake now:

use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

We use a .view call instead of .view_method_call because it's not in the lockup contract. With lockup contract, we can import the function and directly call; but we only have the wasm file of the staking contract which we call from there, so we use view instead to simulate calling near view from the near-cli.

To update the staking pool balance globally, we refresh it. We can simulate the refresh here. Recall that the function is useful when owner wants to receive rewards during unstaking for querying total balance in the pool that could be withdrawn. In reality, this should be "no-op", meaning that the refresh should be done automatically with staking.

use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

Next, we simulate the rewards, update the staking pool status by "pinging" it, then check that our staked amount is now larger than previously deposited (the final line assertion). The ping action has no arguments, and we don't need to define a &json! just for it: we can use b"" which is a bytes with nothing inside, it'll get parsed as a &[u8] type as arguments.

use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

We refresh the staking balance again and check for two things:

  • After refresh, the new stake amount is still larger than previously staked.
  • assert_eq_with_gas will check that the actual amount, after minus gas, is still within what's predictable. This ensures that we don't have a really large difference that comes out of nowhere; and the difference can only be due to gas fee charged.
use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

Then, the process of getting the money back to owner account also has a few step. It start with unstaking the balance. For simplicity, we'll unstake everything. In reality, you can only unstake everything; at least as of writing, the NEAR web wallet only supports unstaking everything staked. It might change in the future.

use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

Given that we have a lot of these view calls and they look the same, we made a small function to deal with that: (actually, it's a closure; or in Python language, a lambda function)

use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

We use a closure because:

  • We cannot clone UserAccount. We just want to borrow it for use and return it after.
  • We want to capture environment because of above case.

Only within a closure you can borrow it temporarily. If we have a separate function, we need to pass the owner object back into the main function after use, otherwise it would meet end of life (lifetime ends).

Note for the unstake, we check for both staked_balance and unstaked_balance, passing different &str to the closure.

Remember we need to wait for some time (4 epochs) before we could withdraw. We also increase some random number (40) to block_height in this process. In reality, block_height increases (approximately) every second because NEAR has (approximately) 1 second finality. 1 epoch is about 12-13 hours; so block height should increase much much more than 40.

use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

Then we unlock and withdraw. In reality, unlocking balance is done automatically; it's just the simulator doesn't have that capability so we manually intervene.

use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

Finally, we unselect the staking pool. That's it, we're done with a normal staking procedure.

use crate::*;


#[test]
fn staking() {
    let lockup_amount = to_yocto("1000");
    let (root, foundation, owner, staking_pool) = basic_setup();

    let staking_account: AccountId = STAKING_POOL_ACCOUNT_ID.parse().unwrap();
    let lockup_account : AccountId = LOCKUP_ACCOUNT_ID.parse().unwrap();

    let lockup = deploy!(
        contract: LockupContractContract,
        contract_id: lockup_account,
        bytes: &LOCKUP_WASM_BYTES,
        signer_account: root,
        deposit: MIN_BALANCE_FOR_STORAGE + lockup_amount,
        gas: MAX_GAS,
        init_method: new(
            owner.account_id.clone(),
            1000000000.into(),
            None,
            TransfersInformation::TransfersDisabled {
                transfer_poll_account_id: "transfer-poll".parse().unwrap(),
            },
            None,
            None,
            STAKING_POOL_WHITELIST_ACCOUNT_ID.parse().unwrap(),
            None
        )
    );

    let get_account_balance = |method: &str| -> U128 {
      owner.view(
        STAKING_POOL_ACCOUNT_ID.parse().unwrap(),
        method,
        &json!({
          "account_id": LOCKUP_ACCOUNT_ID.to_string()
        }).to_string().into_bytes(),
      ).unwrap_json()
    };

    let owner_staking_account = &owner;

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);

    // Selecting staking pool
    owner_staking_account.function_call(
      lockup.contract.select_staking_pool(staking_account.clone()),
      MAX_GAS,
      0,
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id() 
    ).unwrap_json();

    assert_eq!(res, Some(STAKING_POOL_ACCOUNT_ID.parse().unwrap()));

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    // Depositing to the staking pool
    let staking_amount = lockup_amount - to_yocto("100");

    owner_staking_account.function_call(
      lockup.contract.deposit_to_staking_pool(U128(staking_amount)),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, staking_amount);

    // Staking on the staking pool
    owner_staking_account.function_call(
      lockup.contract.stake(U128(staking_amount)), 
      MAX_GAS, 
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_yocto_eq(res.0, staking_amount);

    // Refreshing staking balance. Should be NOOP
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_yocto_eq(res.0, staking_amount);

    // Simulating rewards
    foundation.transfer(
      staking_account.clone(),
      to_yocto("10")
    ).assert_success();

    // Pinging the staking pool
    foundation.call(
      staking_account.clone(),
      "ping",
      b"",
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = get_account_balance("get_account_staked_balance");

    let new_stake_amount = res.0;
    assert!(new_stake_amount > staking_amount);

    // Refresh staking balance again
    owner_staking_account.function_call(
      lockup.contract.refresh_staking_pool_balance(),
      MAX_GAS,
      0
    ).assert_success();

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    let new_total_balance = res.0;
    assert!(new_total_balance >= new_stake_amount);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    // Account for gas rewards
    assert_eq_with_gas(
      res.0, 
      new_total_balance - staking_amount
    );

    // Unstaking everything
    let res: bool = owner_staking_account.function_call(
      lockup.contract.unstake(U128(new_stake_amount)),
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = get_account_balance("get_account_staked_balance");

    assert_eq_with_gas(res.0, 0);

    let res: U128 = get_account_balance("get_account_unstaked_balance");

    assert!(res.0 >= new_total_balance);
    
    root.borrow_runtime_mut().cur_block.block_height += 40;
    root.borrow_runtime_mut().cur_block.epoch_height += 4;

    // The standalone runtime doesn't unlock locked balance. Need to manually intervene.
    let mut pool = staking_pool.account().unwrap();
    pool.amount += pool.locked;
    pool.locked = 0;
    staking_pool
        .borrow_runtime_mut()
        .force_account_update(STAKING_POOL_ACCOUNT_ID.parse().unwrap(), &pool);

    // Withdrawing everything from the staking pool
    let res: bool = owner_staking_account.function_call(
      lockup.contract.withdraw_from_staking_pool(U128(new_total_balance)) ,
      MAX_GAS,
      0
    ).unwrap_json();

    assert!(res);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_known_deposited_balance() 
    ).unwrap_json();

    assert_eq!(res.0, 0);

    let res: U128 = owner.view_method_call(
      lockup.contract.get_owners_balance() 
    ).unwrap_json();

    assert_eq_with_gas(res.0, new_stake_amount - staking_amount);

    // Unselecting the staking pool
    owner_staking_account.function_call(
      lockup.contract.unselect_staking_pool(),
      MAX_GAS,
      0
    ).assert_success();

    let res: Option<AccountId> = owner.view_method_call(
      lockup.contract.get_staking_pool_account_id()
    ).unwrap_json();

    assert_eq!(res, None);
}

There are a few more functions which one won't include them here; but you can always refer to them. Link in References.

There might be some function that one missed or left out because it's repeating. You shall add them back in yourself if you found "no method error". Could you guess the function that we missed out? (We do deliberately miss out one function that you can't see but it's hidden inside). (Remember to rebuild after you change the code!)

That ends what we want to speak about for this chapter. In the next few chapters, we would look into developing our own app. Particularly, something to prepare you for Near Certified Developer (NCD) in Learn Near Club (LNC). (Note the NCD from near.university might be more difficult than NCD from LNC; as far as one is aware based on the syllabus; so be prepared. In NCD for LNC, you demo with other people for 60 minutes; but the website in Near University mentioned you demo yourself for 60 minutes in NCD for Near university).

References