Test for Foundation (Tentative?)

We will write some local test to foundations. To first start, let's know that some of the test utilities are put inside src/tests/test_utils.rs which can be imported. This allows not to cramp everything inside the actual testing module.

We won't be demonstrating this, but you can visit the page here.

One originally want to separate this into the foundation.rs file but found it couldn't be imported in any ways unless create the same file multiple times and put them in different folders, which is just redundant. Hopefully Rust can change this in the future as it makes no sense to import from different test utils if you really pull it out to share them across multiple files; but for now, we'll have to put everything in the lib.rs one afraid.

Anyway, let's take a look at some of the tests related to foundation.rs. We shall start with the test related to terminate_vesting.

Usually when we do test, we try not to write too much tests, else if we change a small thing then the test also changes, which is redundant and unsustainable. But if there is a need to perform unit test on a module, then we can try to group test together; as in some test run multiple functions, and these related functions could be tested together.

Here, terminate_vesting itself can test lots of stuffs, like test for self.assert_called_by_foundation() (though one suggests this as a separate test if it's always being used), test for require!, test for correct output if everything goes correctly, sometimes test for edge cases. From the smart contract, it makes sense that tests should cover two category:

  1. Output is expected. We don't want to have a user transfer 10 NEAR to account B but account B only receive 5 NEAR and the rest burnt unexpectedly.
  2. Security. The assertions are in place to confirm for security. Hence, our test may be designed from the point of view of hackers and check whether our contract panic when somebody tries to hack through the contract.

Of course, we aren't God, hence we won't be able to cover everything. We try to cover as much as possible what our team and the public can think of based on suggestions by others or your own thoughts and put those that are critical in.

Another thing one wants to say is, while the signer_account_id and predecessor_account_id in near_sdk are no longer just pure String (which previously pub type AccountId = String is used in v3.1.0), the near-sdk::VMContext haven't change to reflect that change, so we need to convert our AccountId back .to_string() as a temporary workaround. One is doing this rather than just returning pure String to ease understanding, though of course one agree that this is more redundant than just returning String from functions in terms of optimization purposes.

Another change that exist is context.is_view. You won't see it in the first 3 tests (you'll see it in the last text). Basically when you set this to false, execution will not charge gas fee while false, and continue charging after you change it back to true.

Explaining a bit on the =4.0.0-pre.5 change.

First, the change on current_account_id now taking in AccountId. But bear in mind, one still uses =4.0.0-pre.4 because of this issue. The AccountId here it takes in isn't near-sdk::AccountId, but near-account-id::AccountId (and it has to be near-account-id-0.10.0 version (don't use 0.12.0 version or it says near-account-id::AccountId is not equal near-account-id::AccountId when the words reads the same, but difference comes in version)). So we can't use the latest version just yet until they port the correct AccountId in.

Next is the is_view, which is now a Method, hence you cannot assign value to it; it's read only. However, after reading the docs, we found something similar. This puts the execution in view mode and defines its configuration. The assignment is different here, as it now takes in not a bool but an Option<ViewConfig>. Let's explain.

First, ViewConfig is from near-primitives-core, so you need to change your toml file to import that library:

[dependencies]
near-primitives-core = "0.12.0"

(you could do anything above 0.10.0, as of writing).

Then you can look at the docs and find how to import it, we see that it contains a max_gas_burnt field, which I won't speak what this means so that you can take your effort and look it up yourself.

Finish looking up? Great. Let's continue.

If the max_gas_burnt field is defined, or perhaps it doesn't have to be defined, one don't know at this point (we'll come back to it later), if we have Some(ViewConfig {}) (is this valid?) or Some(ViewConfig { max_gas_burnt: 10u128.pow(12) }), then it's similar to config.is_view = true. Luckily for the basic_context() that we're going to use later, we just need this to be None, so we don't need to care at this point how it works out.

That's it, let's continue.

Continue on with our test cases

Let's know that we did some shortcuts to remove redundant lines of code:

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

So the first test, we shall test calling by non-foundation and expect it to fail.

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

Then we see that we passed in a vesting_schedule_with_salt, and we want to test and see if we change the value inside, it'll fail as we expected it to. Particularly, it'll fail if we use a different salt, and it'll fail if we have a different vesting schedule.

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

Then we'll skip some test like testing for no vesting_schedule_with_salt provided.

Note that we use assert rather than require here. For consistency we could continue to use require; but it is not required. require is a lightweight version of assert and it's useful during deployment, but during testing, you can have heavyweight (well, it's actually not that heavyweight, just that on-chain requires computation that cost gas fee so it's best to decrease this as much as possible) macro.

Next, let's look at termination with staking. This is one of the longest assertion as it is more like a simulation test than normal test. As we know, with staking, we will get "deficit" as staking gets us rewards, hence account balance will increase to more than when we stake (minus gas fee and storage fee). We want to simulate staking and ensure it works.

(The original staking contract might be too long and we might reduce it to smaller tests for easier explanation; but in real, you just write one single test as seen here in line 1728 onwards as of writing (or search for test_termination_with_staking())).

The first part is the setup:

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

context.is_view based on the docs, it says the execution should not charge any costs while it's true.

As of writing this section, we commented out two of the test here (contract.get_owners_balance() and contract.get_liquid_owners_balance()) that should be there to assert the initialization is correct to make things simpler (the functions are getters which depends on other getters which depends on other getters...). You however might not be seeing they covered up because we're writing a single contract, and when we have the function we might unleash them.

The others, as the test mentions, check that we have the correct unvested and vested amount. Those numbers are determined by the vested schedule, and should be such and such at such and such timing.

Next is select staking pool and deposit to staking pool. As usual, we'll skip quite a lot of tests that should be there to make things simpler. (We might not even include those as comments here, so make sure to check back the original code and find out).

There's another thing one wants to test, which are the rest of the foundation, but really they requires the select_staking_pool or it'll fail with staking pool not selected; hence we can only test the fail case not the not-failed case until we worked through the owner.rs. We'll come back to it later.

The failed case are as below. We'll abandon this test for proper testing later in future pages. I'll put the rest of this code here (after changing the name too)!

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

With that, let's move on to Foundation callbacks.