@aptos-labs/ts-sdk 3.0.0 → 3.1.0

package/src/internal/move/jwks/build/jwk/sources/dependencies/AptosFramework/staking_contract.move

@@ -0,0 +1,1618 @@
+/// Allow stakers and operators to enter a staking contract with reward sharing.
+/// The main accounting logic in a staking contract consists of 2 parts:
+/// 1. Tracks how much commission needs to be paid out to the operator. This is tracked with an increasing principal
+/// amount that's updated every time the operator requests commission, the staker withdraws funds, or the staker
+/// switches operators.
+/// 2. Distributions of funds to operators (commissions) and stakers (stake withdrawals) use the shares model provided
+/// by the pool_u64 to track shares that increase in price as the stake pool accumulates rewards.
+///
+/// Example flow:
+/// 1. A staker creates a staking contract with an operator by calling create_staking_contract() with 100 coins of
+/// initial stake and commission = 10%. This means the operator will receive 10% of any accumulated rewards. A new stake
+/// pool will be created and hosted in a separate account that's controlled by the staking contract.
+/// 2. The operator sets up a validator node and, once ready, joins the validator set by calling stake::join_validator_set
+/// 3. After some time, the stake pool gains rewards and now has 150 coins.
+/// 4. Operator can now call request_commission. 10% of (150 - 100) = 5 coins will be unlocked from the stake pool. The
+/// staker's principal is now updated from 100 to 145 (150 coins - 5 coins of commission). The pending distribution pool
+/// has 5 coins total and the operator owns all 5 shares of it.
+/// 5. Some more time has passed. The pool now has 50 more coins in rewards and a total balance of 195. The operator
+/// calls request_commission again. Since the previous 5 coins have now become withdrawable, it'll be deposited into the
+/// operator's account first. Their new commission will be 10% of (195 coins - 145 principal) = 5 coins. Principal is
+/// updated to be 190 (195 - 5). Pending distribution pool has 5 coins and operator owns all 5 shares.
+/// 6. Staker calls unlock_stake to unlock 50 coins of stake, which gets added to the pending distribution pool. Based
+/// on shares math, staker will be owning 50 shares and operator still owns 5 shares of the 55-coin pending distribution
+/// pool.
+/// 7. Some time passes and the 55 coins become fully withdrawable from the stake pool. Due to accumulated rewards, the
+/// 55 coins become 70 coins. Calling distribute() distributes 6 coins to the operator and 64 coins to the validator.
+module aptos_framework::staking_contract {
+    use std::bcs;
+    use std::error;
+    use std::features;
+    use std::signer;
+    use std::vector;
+
+    use aptos_std::pool_u64::{Self, Pool};
+    use aptos_std::simple_map::{Self, SimpleMap};
+
+    use aptos_framework::account::{Self, SignerCapability};
+    use aptos_framework::aptos_account;
+    use aptos_framework::aptos_coin::AptosCoin;
+    use aptos_framework::coin::{Self, Coin};
+    use aptos_framework::event::{EventHandle, emit, emit_event};
+    use aptos_framework::stake::{Self, OwnerCapability};
+    use aptos_framework::staking_config;
+
+    const SALT: vector<u8> = b"aptos_framework::staking_contract";
+
+    /// Store amount must be at least the min stake required for a stake pool to join the validator set.
+    const EINSUFFICIENT_STAKE_AMOUNT: u64 = 1;
+    /// Commission percentage has to be between 0 and 100.
+    const EINVALID_COMMISSION_PERCENTAGE: u64 = 2;
+    /// Staker has no staking contracts.
+    const ENO_STAKING_CONTRACT_FOUND_FOR_STAKER: u64 = 3;
+    /// No staking contract between the staker and operator found.
+    const ENO_STAKING_CONTRACT_FOUND_FOR_OPERATOR: u64 = 4;
+    /// Staking contracts can't be merged.
+    const ECANT_MERGE_STAKING_CONTRACTS: u64 = 5;
+    /// The staking contract already exists and cannot be re-created.
+    const ESTAKING_CONTRACT_ALREADY_EXISTS: u64 = 6;
+    /// Not enough active stake to withdraw. Some stake might still pending and will be active in the next epoch.
+    const EINSUFFICIENT_ACTIVE_STAKE_TO_WITHDRAW: u64 = 7;
+    /// Caller must be either the staker, operator, or beneficiary.
+    const ENOT_STAKER_OR_OPERATOR_OR_BENEFICIARY: u64 = 8;
+    /// Changing beneficiaries for operators is not supported.
+    const EOPERATOR_BENEFICIARY_CHANGE_NOT_SUPPORTED: u64 = 9;
+
+    /// Maximum number of distributions a stake pool can support.
+    const MAXIMUM_PENDING_DISTRIBUTIONS: u64 = 20;
+
+    #[resource_group(scope = module_)]
+    struct StakingGroupContainer {}
+
+    struct StakingContract has store {
+        // Recorded principal after the last commission distribution.
+        // This is only used to calculate the commission the operator should be receiving.
+        principal: u64,
+        pool_address: address,
+        // The stake pool's owner capability. This can be used to control funds in the stake pool.
+        owner_cap: OwnerCapability,
+        commission_percentage: u64,
+        // Current distributions, including operator commission withdrawals and staker's partial withdrawals.
+        distribution_pool: Pool,
+        // Just in case we need the SignerCap for stake pool account in the future.
+        signer_cap: SignerCapability,
+    }
+
+    struct Store has key {
+        staking_contracts: SimpleMap<address, StakingContract>,
+
+        // Events.
+        create_staking_contract_events: EventHandle<CreateStakingContractEvent>,
+        update_voter_events: EventHandle<UpdateVoterEvent>,
+        reset_lockup_events: EventHandle<ResetLockupEvent>,
+        add_stake_events: EventHandle<AddStakeEvent>,
+        request_commission_events: EventHandle<RequestCommissionEvent>,
+        unlock_stake_events: EventHandle<UnlockStakeEvent>,
+        switch_operator_events: EventHandle<SwitchOperatorEvent>,
+        add_distribution_events: EventHandle<AddDistributionEvent>,
+        distribute_events: EventHandle<DistributeEvent>,
+    }
+
+    struct BeneficiaryForOperator has key {
+        beneficiary_for_operator: address,
+    }
+
+    struct UpdateCommissionEvent has drop, store {
+        staker: address,
+        operator: address,
+        old_commission_percentage: u64,
+        new_commission_percentage: u64,
+    }
+
+    #[event]
+    struct UpdateCommission has drop, store {
+        staker: address,
+        operator: address,
+        old_commission_percentage: u64,
+        new_commission_percentage: u64,
+    }
+
+    #[resource_group_member(group = aptos_framework::staking_contract::StakingGroupContainer)]
+    struct StakingGroupUpdateCommissionEvent has key {
+        update_commission_events: EventHandle<UpdateCommissionEvent>,
+    }
+
+    #[event]
+    struct CreateStakingContract has drop, store {
+        operator: address,
+        voter: address,
+        pool_address: address,
+        principal: u64,
+        commission_percentage: u64,
+    }
+
+    #[event]
+    struct UpdateVoter has drop, store {
+        operator: address,
+        pool_address: address,
+        old_voter: address,
+        new_voter: address,
+    }
+
+    #[event]
+    struct ResetLockup has drop, store {
+        operator: address,
+        pool_address: address,
+    }
+
+    #[event]
+    struct AddStake has drop, store {
+        operator: address,
+        pool_address: address,
+        amount: u64
+    }
+
+    #[event]
+    struct RequestCommission has drop, store {
+        operator: address,
+        pool_address: address,
+        accumulated_rewards: u64,
+        commission_amount: u64,
+    }
+
+    #[event]
+    struct UnlockStake has drop, store {
+        operator: address,
+        pool_address: address,
+        amount: u64,
+        commission_paid: u64,
+    }
+
+    #[event]
+    struct SwitchOperator has drop, store {
+        old_operator: address,
+        new_operator: address,
+        pool_address: address,
+    }
+
+    #[event]
+    struct AddDistribution has drop, store {
+        operator: address,
+        pool_address: address,
+        amount: u64,
+    }
+
+    #[event]
+    struct Distribute has drop, store {
+        operator: address,
+        pool_address: address,
+        recipient: address,
+        amount: u64,
+    }
+
+    #[event]
+    struct SetBeneficiaryForOperator has drop, store {
+        operator: address,
+        old_beneficiary: address,
+        new_beneficiary: address,
+    }
+
+    struct CreateStakingContractEvent has drop, store {
+        operator: address,
+        voter: address,
+        pool_address: address,
+        principal: u64,
+        commission_percentage: u64,
+    }
+
+    struct UpdateVoterEvent has drop, store {
+        operator: address,
+        pool_address: address,
+        old_voter: address,
+        new_voter: address,
+    }
+
+    struct ResetLockupEvent has drop, store {
+        operator: address,
+        pool_address: address,
+    }
+
+    struct AddStakeEvent has drop, store {
+        operator: address,
+        pool_address: address,
+        amount: u64
+    }
+
+    struct RequestCommissionEvent has drop, store {
+        operator: address,
+        pool_address: address,
+        accumulated_rewards: u64,
+        commission_amount: u64,
+    }
+
+    struct UnlockStakeEvent has drop, store {
+        operator: address,
+        pool_address: address,
+        amount: u64,
+        commission_paid: u64,
+    }
+
+    struct SwitchOperatorEvent has drop, store {
+        old_operator: address,
+        new_operator: address,
+        pool_address: address,
+    }
+
+    struct AddDistributionEvent has drop, store {
+        operator: address,
+        pool_address: address,
+        amount: u64,
+    }
+
+    struct DistributeEvent has drop, store {
+        operator: address,
+        pool_address: address,
+        recipient: address,
+        amount: u64,
+    }
+
+    #[view]
+    /// Return the address of the underlying stake pool for the staking contract between the provided staker and
+    /// operator.
+    ///
+    /// This errors out the staking contract with the provided staker and operator doesn't exist.
+    public fun stake_pool_address(staker: address, operator: address): address acquires Store {
+        assert_staking_contract_exists(staker, operator);
+        let staking_contracts = &borrow_global<Store>(staker).staking_contracts;
+        simple_map::borrow(staking_contracts, &operator).pool_address
+    }
+
+    #[view]
+    /// Return the last recorded principal (the amount that 100% belongs to the staker with commission already paid for)
+    /// for staking contract between the provided staker and operator.
+    ///
+    /// This errors out the staking contract with the provided staker and operator doesn't exist.
+    public fun last_recorded_principal(staker: address, operator: address): u64 acquires Store {
+        assert_staking_contract_exists(staker, operator);
+        let staking_contracts = &borrow_global<Store>(staker).staking_contracts;
+        simple_map::borrow(staking_contracts, &operator).principal
+    }
+
+    #[view]
+    /// Return percentage of accumulated rewards that will be paid to the operator as commission for staking contract
+    /// between the provided staker and operator.
+    ///
+    /// This errors out the staking contract with the provided staker and operator doesn't exist.
+    public fun commission_percentage(staker: address, operator: address): u64 acquires Store {
+        assert_staking_contract_exists(staker, operator);
+        let staking_contracts = &borrow_global<Store>(staker).staking_contracts;
+        simple_map::borrow(staking_contracts, &operator).commission_percentage
+    }
+
+    #[view]
+    /// Return a tuple of three numbers:
+    /// 1. The total active stake in the underlying stake pool
+    /// 2. The total accumulated rewards that haven't had commission paid out
+    /// 3. The commission amount owned from those accumulated rewards.
+    ///
+    /// This errors out the staking contract with the provided staker and operator doesn't exist.
+    public fun staking_contract_amounts(staker: address, operator: address): (u64, u64, u64) acquires Store {
+        assert_staking_contract_exists(staker, operator);
+        let staking_contracts = &borrow_global<Store>(staker).staking_contracts;
+        let staking_contract = simple_map::borrow(staking_contracts, &operator);
+        get_staking_contract_amounts_internal(staking_contract)
+    }
+
+    #[view]
+    /// Return the number of pending distributions (e.g. commission, withdrawals from stakers).
+    ///
+    /// This errors out the staking contract with the provided staker and operator doesn't exist.
+    public fun pending_distribution_counts(staker: address, operator: address): u64 acquires Store {
+        assert_staking_contract_exists(staker, operator);
+        let staking_contracts = &borrow_global<Store>(staker).staking_contracts;
+        pool_u64::shareholders_count(&simple_map::borrow(staking_contracts, &operator).distribution_pool)
+    }
+
+    #[view]
+    /// Return true if the staking contract between the provided staker and operator exists.
+    public fun staking_contract_exists(staker: address, operator: address): bool acquires Store {
+        if (!exists<Store>(staker)) {
+            return false
+        };
+
+        let store = borrow_global<Store>(staker);
+        simple_map::contains_key(&store.staking_contracts, &operator)
+    }
+
+    #[view]
+    /// Return the beneficiary address of the operator.
+    public fun beneficiary_for_operator(operator: address): address acquires BeneficiaryForOperator {
+        if (exists<BeneficiaryForOperator>(operator)) {
+            return borrow_global<BeneficiaryForOperator>(operator).beneficiary_for_operator
+        } else {
+            operator
+        }
+    }
+
+    #[view]
+    /// Return the address of the stake pool to be created with the provided staker, operator and seed.
+    public fun get_expected_stake_pool_address(
+        staker: address,
+        operator: address,
+        contract_creation_seed: vector<u8>,
+    ): address {
+        let seed = create_resource_account_seed(staker, operator, contract_creation_seed);
+        account::create_resource_address(&staker, seed)
+    }
+
+    /// Staker can call this function to create a simple staking contract with a specified operator.
+    public entry fun create_staking_contract(
+        staker: &signer,
+        operator: address,
+        voter: address,
+        amount: u64,
+        commission_percentage: u64,
+        // Optional seed used when creating the staking contract account.
+        contract_creation_seed: vector<u8>,
+    ) acquires Store {
+        let staked_coins = coin::withdraw<AptosCoin>(staker, amount);
+        create_staking_contract_with_coins(
+            staker, operator, voter, staked_coins, commission_percentage, contract_creation_seed);
+    }
+
+    /// Staker can call this function to create a simple staking contract with a specified operator.
+    public fun create_staking_contract_with_coins(
+        staker: &signer,
+        operator: address,
+        voter: address,
+        coins: Coin<AptosCoin>,
+        commission_percentage: u64,
+        // Optional seed used when creating the staking contract account.
+        contract_creation_seed: vector<u8>,
+    ): address acquires Store {
+        assert!(
+            commission_percentage >= 0 && commission_percentage <= 100,
+            error::invalid_argument(EINVALID_COMMISSION_PERCENTAGE),
+        );
+        // The amount should be at least the min_stake_required, so the stake pool will be eligible to join the
+        // validator set.
+        let (min_stake_required, _) = staking_config::get_required_stake(&staking_config::get());
+        let principal = coin::value(&coins);
+        assert!(principal >= min_stake_required, error::invalid_argument(EINSUFFICIENT_STAKE_AMOUNT));
+
+        // Initialize Store resource if this is the first time the staker has delegated to anyone.
+        let staker_address = signer::address_of(staker);
+        if (!exists<Store>(staker_address)) {
+            move_to(staker, new_staking_contracts_holder(staker));
+        };
+
+        // Cannot create the staking contract if it already exists.
+        let store = borrow_global_mut<Store>(staker_address);
+        let staking_contracts = &mut store.staking_contracts;
+        assert!(
+            !simple_map::contains_key(staking_contracts, &operator),
+            error::already_exists(ESTAKING_CONTRACT_ALREADY_EXISTS)
+        );
+
+        // Initialize the stake pool in a new resource account. This allows the same staker to contract with multiple
+        // different operators.
+        let (stake_pool_signer, stake_pool_signer_cap, owner_cap) =
+            create_stake_pool(staker, operator, voter, contract_creation_seed);
+
+        // Add the stake to the stake pool.
+        stake::add_stake_with_cap(&owner_cap, coins);
+
+        // Create the contract record.
+        let pool_address = signer::address_of(&stake_pool_signer);
+        simple_map::add(staking_contracts, operator, StakingContract {
+            principal,
+            pool_address,
+            owner_cap,
+            commission_percentage,
+            // Make sure we don't have too many pending recipients in the distribution pool.
+            // Otherwise, a griefing attack is possible where the staker can keep switching operators and create too
+            // many pending distributions. This can lead to out-of-gas failure whenever distribute() is called.
+            distribution_pool: pool_u64::create(MAXIMUM_PENDING_DISTRIBUTIONS),
+            signer_cap: stake_pool_signer_cap,
+        });
+
+        if (std::features::module_event_migration_enabled()) {
+            emit(CreateStakingContract { operator, voter, pool_address, principal, commission_percentage });
+        };
+        emit_event(
+            &mut store.create_staking_contract_events,
+            CreateStakingContractEvent { operator, voter, pool_address, principal, commission_percentage },
+        );
+        pool_address
+    }
+
+    /// Add more stake to an existing staking contract.
+    public entry fun add_stake(staker: &signer, operator: address, amount: u64) acquires Store {
+        let staker_address = signer::address_of(staker);
+        assert_staking_contract_exists(staker_address, operator);
+
+        let store = borrow_global_mut<Store>(staker_address);
+        let staking_contract = simple_map::borrow_mut(&mut store.staking_contracts, &operator);
+
+        // Add the stake to the stake pool.
+        let staked_coins = coin::withdraw<AptosCoin>(staker, amount);
+        stake::add_stake_with_cap(&staking_contract.owner_cap, staked_coins);
+
+        staking_contract.principal = staking_contract.principal + amount;
+        let pool_address = staking_contract.pool_address;
+        if (std::features::module_event_migration_enabled()) {
+            emit(AddStake { operator, pool_address, amount });
+        };
+        emit_event(
+            &mut store.add_stake_events,
+            AddStakeEvent { operator, pool_address, amount },
+        );
+    }
+
+    /// Convenient function to allow the staker to update the voter address in a staking contract they made.
+    public entry fun update_voter(staker: &signer, operator: address, new_voter: address) acquires Store {
+        let staker_address = signer::address_of(staker);
+        assert_staking_contract_exists(staker_address, operator);
+
+        let store = borrow_global_mut<Store>(staker_address);
+        let staking_contract = simple_map::borrow_mut(&mut store.staking_contracts, &operator);
+        let pool_address = staking_contract.pool_address;
+        let old_voter = stake::get_delegated_voter(pool_address);
+        stake::set_delegated_voter_with_cap(&staking_contract.owner_cap, new_voter);
+
+        if (std::features::module_event_migration_enabled()) {
+            emit(UpdateVoter { operator, pool_address, old_voter, new_voter });
+        };
+        emit_event(
+            &mut store.update_voter_events,
+            UpdateVoterEvent { operator, pool_address, old_voter, new_voter },
+        );
+
+    }
+
+    /// Convenient function to allow the staker to reset their stake pool's lockup period to start now.
+    public entry fun reset_lockup(staker: &signer, operator: address) acquires Store {
+        let staker_address = signer::address_of(staker);
+        assert_staking_contract_exists(staker_address, operator);
+
+        let store = borrow_global_mut<Store>(staker_address);
+        let staking_contract = simple_map::borrow_mut(&mut store.staking_contracts, &operator);
+        let pool_address = staking_contract.pool_address;
+        stake::increase_lockup_with_cap(&staking_contract.owner_cap);
+
+        if (std::features::module_event_migration_enabled()) {
+            emit(ResetLockup { operator, pool_address });
+        };
+        emit_event(&mut store.reset_lockup_events, ResetLockupEvent { operator, pool_address });
+    }
+
+    /// Convenience function to allow a staker to update the commission percentage paid to the operator.
+    /// TODO: fix the typo in function name. commision -> commission
+    public entry fun update_commision(
+        staker: &signer,
+        operator: address,
+        new_commission_percentage: u64
+    ) acquires Store, BeneficiaryForOperator, StakingGroupUpdateCommissionEvent {
+        assert!(
+            new_commission_percentage >= 0 && new_commission_percentage <= 100,
+            error::invalid_argument(EINVALID_COMMISSION_PERCENTAGE),
+        );
+
+        let staker_address = signer::address_of(staker);
+        assert!(exists<Store>(staker_address), error::not_found(ENO_STAKING_CONTRACT_FOUND_FOR_STAKER));
+
+        let store = borrow_global_mut<Store>(staker_address);
+        let staking_contract = simple_map::borrow_mut(&mut store.staking_contracts, &operator);
+        distribute_internal(staker_address, operator, staking_contract, &mut store.distribute_events);
+        request_commission_internal(
+            operator,
+            staking_contract,
+            &mut store.add_distribution_events,
+            &mut store.request_commission_events,
+        );
+        let old_commission_percentage = staking_contract.commission_percentage;
+        staking_contract.commission_percentage = new_commission_percentage;
+        if (!exists<StakingGroupUpdateCommissionEvent>(staker_address)) {
+            move_to(
+                staker,
+                StakingGroupUpdateCommissionEvent {
+                    update_commission_events: account::new_event_handle<UpdateCommissionEvent>(
+                        staker
+                    )
+                }
+            )
+        };
+        if (std::features::module_event_migration_enabled()) {
+            emit(
+                UpdateCommission { staker: staker_address, operator, old_commission_percentage, new_commission_percentage }
+            );
+        };
+        emit_event(
+            &mut borrow_global_mut<StakingGroupUpdateCommissionEvent>(staker_address).update_commission_events,
+            UpdateCommissionEvent { staker: staker_address, operator, old_commission_percentage, new_commission_percentage }
+        );
+    }
+
+    /// Unlock commission amount from the stake pool. Operator needs to wait for the amount to become withdrawable
+    /// at the end of the stake pool's lockup period before they can actually can withdraw_commission.
+    ///
+    /// Only staker, operator or beneficiary can call this.
+    public entry fun request_commission(
+        account: &signer,
+        staker: address,
+        operator: address
+    ) acquires Store, BeneficiaryForOperator {
+        let account_addr = signer::address_of(account);
+        assert!(
+            account_addr == staker || account_addr == operator || account_addr == beneficiary_for_operator(operator),
+            error::unauthenticated(ENOT_STAKER_OR_OPERATOR_OR_BENEFICIARY)
+        );
+        assert_staking_contract_exists(staker, operator);
+
+        let store = borrow_global_mut<Store>(staker);
+        let staking_contract = simple_map::borrow_mut(&mut store.staking_contracts, &operator);
+        // Short-circuit if zero commission.
+        if (staking_contract.commission_percentage == 0) {
+            return
+        };
+
+        // Force distribution of any already inactive stake.
+        distribute_internal(staker, operator, staking_contract, &mut store.distribute_events);
+
+        request_commission_internal(
+            operator,
+            staking_contract,
+            &mut store.add_distribution_events,
+            &mut store.request_commission_events,
+        );
+    }
+
+    fun request_commission_internal(
+        operator: address,
+        staking_contract: &mut StakingContract,
+        add_distribution_events: &mut EventHandle<AddDistributionEvent>,
+        request_commission_events: &mut EventHandle<RequestCommissionEvent>,
+    ): u64 {
+        // Unlock just the commission portion from the stake pool.
+        let (total_active_stake, accumulated_rewards, commission_amount) =
+            get_staking_contract_amounts_internal(staking_contract);
+        staking_contract.principal = total_active_stake - commission_amount;
+
+        // Short-circuit if there's no commission to pay.
+        if (commission_amount == 0) {
+            return 0
+        };
+
+        // Add a distribution for the operator.
+        add_distribution(operator, staking_contract, operator, commission_amount, add_distribution_events);
+
+        // Request to unlock the commission from the stake pool.
+        // This won't become fully unlocked until the stake pool's lockup expires.
+        stake::unlock_with_cap(commission_amount, &staking_contract.owner_cap);
+
+        let pool_address = staking_contract.pool_address;
+        if (std::features::module_event_migration_enabled()) {
+            emit(RequestCommission { operator, pool_address, accumulated_rewards, commission_amount });
+        };
+        emit_event(
+            request_commission_events,
+            RequestCommissionEvent { operator, pool_address, accumulated_rewards, commission_amount },
+        );
+
+        commission_amount
+    }
+
+    /// Staker can call this to request withdrawal of part or all of their staking_contract.
+    /// This also triggers paying commission to the operator for accounting simplicity.
+    public entry fun unlock_stake(
+        staker: &signer,
+        operator: address,
+        amount: u64
+    ) acquires Store, BeneficiaryForOperator {
+        // Short-circuit if amount is 0.
+        if (amount == 0) return;
+
+        let staker_address = signer::address_of(staker);
+        assert_staking_contract_exists(staker_address, operator);
+
+        let store = borrow_global_mut<Store>(staker_address);
+        let staking_contract = simple_map::borrow_mut(&mut store.staking_contracts, &operator);
+
+        // Force distribution of any already inactive stake.
+        distribute_internal(staker_address, operator, staking_contract, &mut store.distribute_events);
+
+        // For simplicity, we request commission to be paid out first. This avoids having to ensure to staker doesn't
+        // withdraw into the commission portion.
+        let commission_paid = request_commission_internal(
+            operator,
+            staking_contract,
+            &mut store.add_distribution_events,
+            &mut store.request_commission_events,
+        );
+
+        // If there's less active stake remaining than the amount requested (potentially due to commission),
+        // only withdraw up to the active amount.
+        let (active, _, _, _) = stake::get_stake(staking_contract.pool_address);
+        if (active < amount) {
+            amount = active;
+        };
+        staking_contract.principal = staking_contract.principal - amount;
+
+        // Record a distribution for the staker.
+        add_distribution(operator, staking_contract, staker_address, amount, &mut store.add_distribution_events);
+
+        // Request to unlock the distribution amount from the stake pool.
+        // This won't become fully unlocked until the stake pool's lockup expires.
+        stake::unlock_with_cap(amount, &staking_contract.owner_cap);
+
+        let pool_address = staking_contract.pool_address;
+        if (std::features::module_event_migration_enabled()) {
+            emit(UnlockStake { pool_address, operator, amount, commission_paid });
+        };
+        emit_event(
+            &mut store.unlock_stake_events,
+            UnlockStakeEvent { pool_address, operator, amount, commission_paid },
+        );
+    }
+
+    /// Unlock all accumulated rewards since the last recorded principals.
+    public entry fun unlock_rewards(staker: &signer, operator: address) acquires Store, BeneficiaryForOperator {
+        let staker_address = signer::address_of(staker);
+        assert_staking_contract_exists(staker_address, operator);
+
+        // Calculate how much rewards belongs to the staker after commission is paid.
+        let (_, accumulated_rewards, unpaid_commission) = staking_contract_amounts(staker_address, operator);
+        let staker_rewards = accumulated_rewards - unpaid_commission;
+        unlock_stake(staker, operator, staker_rewards);
+    }
+
+    /// Allows staker to switch operator without going through the lenghthy process to unstake, without resetting commission.
+    public entry fun switch_operator_with_same_commission(
+        staker: &signer,
+        old_operator: address,
+        new_operator: address,
+    ) acquires Store, BeneficiaryForOperator {
+        let staker_address = signer::address_of(staker);
+        assert_staking_contract_exists(staker_address, old_operator);
+
+        let commission_percentage = commission_percentage(staker_address, old_operator);
+        switch_operator(staker, old_operator, new_operator, commission_percentage);
+    }
+
+    /// Allows staker to switch operator without going through the lenghthy process to unstake.
+    public entry fun switch_operator(
+        staker: &signer,
+        old_operator: address,
+        new_operator: address,
+        new_commission_percentage: u64,
+    ) acquires Store, BeneficiaryForOperator {
+        let staker_address = signer::address_of(staker);
+        assert_staking_contract_exists(staker_address, old_operator);
+
+        // Merging two existing staking contracts is too complex as we'd need to merge two separate stake pools.
+        let store = borrow_global_mut<Store>(staker_address);
+        let staking_contracts = &mut store.staking_contracts;
+        assert!(
+            !simple_map::contains_key(staking_contracts, &new_operator),
+            error::invalid_state(ECANT_MERGE_STAKING_CONTRACTS),
+        );
+
+        let (_, staking_contract) = simple_map::remove(staking_contracts, &old_operator);
+        // Force distribution of any already inactive stake.
+        distribute_internal(staker_address, old_operator, &mut staking_contract, &mut store.distribute_events);
+
+        // For simplicity, we request commission to be paid out first. This avoids having to ensure to staker doesn't
+        // withdraw into the commission portion.
+        request_commission_internal(
+            old_operator,
+            &mut staking_contract,
+            &mut store.add_distribution_events,
+            &mut store.request_commission_events,
+        );
+
+        // Update the staking contract's commission rate and stake pool's operator.
+        stake::set_operator_with_cap(&staking_contract.owner_cap, new_operator);
+        staking_contract.commission_percentage = new_commission_percentage;
+
+        let pool_address = staking_contract.pool_address;
+        simple_map::add(staking_contracts, new_operator, staking_contract);
+        if (std::features::module_event_migration_enabled()) {
+            emit(SwitchOperator { pool_address, old_operator, new_operator });
+        };
+        emit_event(
+            &mut store.switch_operator_events,
+            SwitchOperatorEvent { pool_address, old_operator, new_operator }
+        );
+    }
+
+    /// Allows an operator to change its beneficiary. Any existing unpaid commission rewards will be paid to the new
+    /// beneficiary. To ensures payment to the current beneficiary, one should first call `distribute` before switching
+    /// the beneficiary. An operator can set one beneficiary for staking contract pools, not a separate one for each pool.
+    public entry fun set_beneficiary_for_operator(
+        operator: &signer,
+        new_beneficiary: address
+    ) acquires BeneficiaryForOperator {
+        assert!(features::operator_beneficiary_change_enabled(), std::error::invalid_state(
+            EOPERATOR_BENEFICIARY_CHANGE_NOT_SUPPORTED
+        ));
+        // The beneficiay address of an operator is stored under the operator's address.
+        // So, the operator does not need to be validated with respect to a staking pool.
+        let operator_addr = signer::address_of(operator);
+        let old_beneficiary = beneficiary_for_operator(operator_addr);
+        if (exists<BeneficiaryForOperator>(operator_addr)) {
+            borrow_global_mut<BeneficiaryForOperator>(operator_addr).beneficiary_for_operator = new_beneficiary;
+        } else {
+            move_to(operator, BeneficiaryForOperator { beneficiary_for_operator: new_beneficiary });
+        };
+
+        emit(SetBeneficiaryForOperator {
+            operator: operator_addr,
+            old_beneficiary,
+            new_beneficiary,
+        });
+    }
+
+    /// Allow anyone to distribute already unlocked funds. This does not affect reward compounding and therefore does
+    /// not need to be restricted to just the staker or operator.
+    public entry fun distribute(staker: address, operator: address) acquires Store, BeneficiaryForOperator {
+        assert_staking_contract_exists(staker, operator);
+        let store = borrow_global_mut<Store>(staker);
+        let staking_contract = simple_map::borrow_mut(&mut store.staking_contracts, &operator);
+        distribute_internal(staker, operator, staking_contract, &mut store.distribute_events);
+    }
+
+    /// Distribute all unlocked (inactive) funds according to distribution shares.
+    fun distribute_internal(
+        staker: address,
+        operator: address,
+        staking_contract: &mut StakingContract,
+        distribute_events: &mut EventHandle<DistributeEvent>,
+    ) acquires BeneficiaryForOperator {
+        let pool_address = staking_contract.pool_address;
+        let (_, inactive, _, pending_inactive) = stake::get_stake(pool_address);
+        let total_potential_withdrawable = inactive + pending_inactive;
+        let coins = stake::withdraw_with_cap(&staking_contract.owner_cap, total_potential_withdrawable);
+        let distribution_amount = coin::value(&coins);
+        if (distribution_amount == 0) {
+            coin::destroy_zero(coins);
+            return
+        };
+
+        let distribution_pool = &mut staking_contract.distribution_pool;
+        update_distribution_pool(
+            distribution_pool, distribution_amount, operator, staking_contract.commission_percentage);
+
+        // Buy all recipients out of the distribution pool.
+        while (pool_u64::shareholders_count(distribution_pool) > 0) {
+            let recipients = pool_u64::shareholders(distribution_pool);
+            let recipient = *vector::borrow(&mut recipients, 0);
+            let current_shares = pool_u64::shares(distribution_pool, recipient);
+            let amount_to_distribute = pool_u64::redeem_shares(distribution_pool, recipient, current_shares);
+            // If the recipient is the operator, send the commission to the beneficiary instead.
+            if (recipient == operator) {
+                recipient = beneficiary_for_operator(operator);
+            };
+            aptos_account::deposit_coins(recipient, coin::extract(&mut coins, amount_to_distribute));
+
+            if (std::features::module_event_migration_enabled()) {
+                emit(Distribute { operator, pool_address, recipient, amount: amount_to_distribute });
+            };
+            emit_event(
+                distribute_events,
+                DistributeEvent { operator, pool_address, recipient, amount: amount_to_distribute }
+            );
+        };
+
+        // In case there's any dust left, send them all to the staker.
+        if (coin::value(&coins) > 0) {
+            aptos_account::deposit_coins(staker, coins);
+            pool_u64::update_total_coins(distribution_pool, 0);
+        } else {
+            coin::destroy_zero(coins);
+        }
+    }
+
+    /// Assert that a staking_contract exists for the staker/operator pair.
+    fun assert_staking_contract_exists(staker: address, operator: address) acquires Store {
+        assert!(exists<Store>(staker), error::not_found(ENO_STAKING_CONTRACT_FOUND_FOR_STAKER));
+        let staking_contracts = &mut borrow_global_mut<Store>(staker).staking_contracts;
+        assert!(
+            simple_map::contains_key(staking_contracts, &operator),
+            error::not_found(ENO_STAKING_CONTRACT_FOUND_FOR_OPERATOR),
+        );
+    }
+
+    /// Add a new distribution for `recipient` and `amount` to the staking contract's distributions list.
+    fun add_distribution(
+        operator: address,
+        staking_contract: &mut StakingContract,
+        recipient: address,
+        coins_amount: u64,
+        add_distribution_events: &mut EventHandle<AddDistributionEvent>
+    ) {
+        let distribution_pool = &mut staking_contract.distribution_pool;
+        let (_, _, _, total_distribution_amount) = stake::get_stake(staking_contract.pool_address);
+        update_distribution_pool(
+            distribution_pool, total_distribution_amount, operator, staking_contract.commission_percentage);
+
+        pool_u64::buy_in(distribution_pool, recipient, coins_amount);
+        let pool_address = staking_contract.pool_address;
+        if (std::features::module_event_migration_enabled()) {
+            emit(AddDistribution { operator, pool_address, amount: coins_amount });
+        };
+        emit_event(
+            add_distribution_events,
+            AddDistributionEvent { operator, pool_address, amount: coins_amount }
+        );
+    }
+
+    /// Calculate accumulated rewards and commissions since last update.
+    fun get_staking_contract_amounts_internal(staking_contract: &StakingContract): (u64, u64, u64) {
+        // Pending_inactive is not included in the calculation because pending_inactive can only come from:
+        // 1. Outgoing commissions. This means commission has already been extracted.
+        // 2. Stake withdrawals from stakers. This also means commission has already been extracted as
+        // request_commission_internal is called in unlock_stake
+        let (active, _, pending_active, _) = stake::get_stake(staking_contract.pool_address);
+        let total_active_stake = active + pending_active;
+        let accumulated_rewards = total_active_stake - staking_contract.principal;
+        let commission_amount = accumulated_rewards * staking_contract.commission_percentage / 100;
+
+        (total_active_stake, accumulated_rewards, commission_amount)
+    }
+
+    fun create_stake_pool(
+        staker: &signer,
+        operator: address,
+        voter: address,
+        contract_creation_seed: vector<u8>,
+    ): (signer, SignerCapability, OwnerCapability) {
+        // Generate a seed that will be used to create the resource account that hosts the staking contract.
+        let seed = create_resource_account_seed(
+            signer::address_of(staker), operator, contract_creation_seed);
+
+        let (stake_pool_signer, stake_pool_signer_cap) = account::create_resource_account(staker, seed);
+        stake::initialize_stake_owner(&stake_pool_signer, 0, operator, voter);
+
+        // Extract owner_cap from the StakePool, so we have control over it in the staking_contracts flow.
+        // This is stored as part of the staking_contract. Thus, the staker would not have direct control over it without
+        // going through well-defined functions in this module.
+        let owner_cap = stake::extract_owner_cap(&stake_pool_signer);
+
+        (stake_pool_signer, stake_pool_signer_cap, owner_cap)
+    }
+
+    fun update_distribution_pool(
+        distribution_pool: &mut Pool,
+        updated_total_coins: u64,
+        operator: address,
+        commission_percentage: u64,
+    ) {
+        // Short-circuit and do nothing if the pool's total value has not changed.
+        if (pool_u64::total_coins(distribution_pool) == updated_total_coins) {
+            return
+        };
+
+        // Charge all stakeholders (except for the operator themselves) commission on any rewards earnt relatively to the
+        // previous value of the distribution pool.
+        let shareholders = &pool_u64::shareholders(distribution_pool);
+        vector::for_each_ref(shareholders, |shareholder| {
+            let shareholder: address = *shareholder;
+            if (shareholder != operator) {
+                let shares = pool_u64::shares(distribution_pool, shareholder);
+                let previous_worth = pool_u64::balance(distribution_pool, shareholder);
+                let current_worth = pool_u64::shares_to_amount_with_total_coins(
+                    distribution_pool, shares, updated_total_coins);
+                let unpaid_commission = (current_worth - previous_worth) * commission_percentage / 100;
+                // Transfer shares from current shareholder to the operator as payment.
+                // The value of the shares should use the updated pool's total value.
+                let shares_to_transfer = pool_u64::amount_to_shares_with_total_coins(
+                    distribution_pool, unpaid_commission, updated_total_coins);
+                pool_u64::transfer_shares(distribution_pool, shareholder, operator, shares_to_transfer);
+            };
+        });
+
+        pool_u64::update_total_coins(distribution_pool, updated_total_coins);
+    }
+
+    /// Create the seed to derive the resource account address.
+    fun create_resource_account_seed(
+        staker: address,
+        operator: address,
+        contract_creation_seed: vector<u8>,
+    ): vector<u8> {
+        let seed = bcs::to_bytes(&staker);
+        vector::append(&mut seed, bcs::to_bytes(&operator));
+        // Include a salt to avoid conflicts with any other modules out there that might also generate
+        // deterministic resource accounts for the same staker + operator addresses.
+        vector::append(&mut seed, SALT);
+        // Add an extra salt given by the staker in case an account with the same address has already been created.
+        vector::append(&mut seed, contract_creation_seed);
+        seed
+    }
+
+    /// Create a new staking_contracts resource.
+    fun new_staking_contracts_holder(staker: &signer): Store {
+        Store {
+            staking_contracts: simple_map::create<address, StakingContract>(),
+            // Events.
+            create_staking_contract_events: account::new_event_handle<CreateStakingContractEvent>(staker),
+            update_voter_events: account::new_event_handle<UpdateVoterEvent>(staker),
+            reset_lockup_events: account::new_event_handle<ResetLockupEvent>(staker),
+            add_stake_events: account::new_event_handle<AddStakeEvent>(staker),
+            request_commission_events: account::new_event_handle<RequestCommissionEvent>(staker),
+            unlock_stake_events: account::new_event_handle<UnlockStakeEvent>(staker),
+            switch_operator_events: account::new_event_handle<SwitchOperatorEvent>(staker),
+            add_distribution_events: account::new_event_handle<AddDistributionEvent>(staker),
+            distribute_events: account::new_event_handle<DistributeEvent>(staker),
+        }
+    }
+
+    #[test_only]
+    const VALIDATOR_STATUS_ACTIVE: u64 = 2;
+    #[test_only]
+    const VALIDATOR_STATUS_INACTIVE: u64 = 4;
+
+    #[test_only]
+    use aptos_framework::stake::with_rewards;
+
+    #[test_only]
+    const INITIAL_BALANCE: u64 = 100000000000000; // 1M APT coins with 8 decimals.
+
+    #[test_only]
+    const MAXIMUM_STAKE: u64 = 100000000000000000; // 1B APT coins with 8 decimals.
+
+    #[test_only]
+    const MODULE_EVENT: u64 = 26;
+
+    #[test_only]
+    const OPERATOR_BENEFICIARY_CHANGE: u64 = 39;
+
+    #[test_only]
+    public fun setup(aptos_framework: &signer, staker: &signer, operator: &signer, initial_balance: u64) {
+        // Reward rate of 0.1% per epoch.
+        stake::initialize_for_test_custom(
+            aptos_framework,
+            INITIAL_BALANCE,
+            MAXIMUM_STAKE,
+            3600,
+            true,
+            10,
+            10000,
+            1000000
+        );
+
+        let staker_address = signer::address_of(staker);
+        if (!account::exists_at(staker_address)) {
+            account::create_account_for_test(staker_address);
+        };
+        let operator_address = signer::address_of(operator);
+        if (!account::exists_at(operator_address)) {
+            account::create_account_for_test(operator_address);
+        };
+        stake::mint(staker, initial_balance);
+        stake::mint(operator, initial_balance);
+    }
+
+    #[test_only]
+    public fun setup_staking_contract(
+        aptos_framework: &signer,
+        staker: &signer,
+        operator: &signer,
+        amount: u64,
+        commission: u64,
+    ) acquires Store {
+        setup(aptos_framework, staker, operator, amount);
+        let operator_address = signer::address_of(operator);
+
+        // Voter is initially set to operator but then updated to be staker.
+        create_staking_contract(staker, operator_address, operator_address, amount, commission, vector::empty<u8>());
+        std::features::change_feature_flags_for_testing(aptos_framework, vector[MODULE_EVENT, OPERATOR_BENEFICIARY_CHANGE], vector[]);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
+    public entry fun test_end_to_end(
+        aptos_framework: &signer,
+        staker: &signer,
+        operator: &signer
+    ) acquires Store, BeneficiaryForOperator {
+        setup_staking_contract(aptos_framework, staker, operator, INITIAL_BALANCE, 10);
+        let staker_address = signer::address_of(staker);
+        let operator_address = signer::address_of(operator);
+        assert_staking_contract_exists(staker_address, operator_address);
+        assert_staking_contract(staker_address, operator_address, INITIAL_BALANCE, 10);
+
+        // Verify that the stake pool has been set up properly.
+        let pool_address = stake_pool_address(staker_address, operator_address);
+        stake::assert_stake_pool(pool_address, INITIAL_BALANCE, 0, 0, 0);
+        assert!(last_recorded_principal(staker_address, operator_address) == INITIAL_BALANCE, 0);
+
+        // Operator joins the validator set.
+        let (_sk, pk, pop) = stake::generate_identity();
+        stake::join_validator_set_for_test(&pk, &pop, operator, pool_address, true);
+        assert!(stake::get_validator_state(pool_address) == VALIDATOR_STATUS_ACTIVE, 1);
+
+        // Fast forward to generate rewards.
+        stake::end_epoch();
+        let new_balance = with_rewards(INITIAL_BALANCE);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, 0);
+
+        // Operator claims 10% of rewards so far as commissions.
+        let expected_commission_1 = (new_balance - last_recorded_principal(staker_address, operator_address)) / 10;
+        new_balance = new_balance - expected_commission_1;
+        request_commission(operator, staker_address, operator_address);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, expected_commission_1);
+        assert!(last_recorded_principal(staker_address, operator_address) == new_balance, 0);
+        assert_distribution(staker_address, operator_address, operator_address, expected_commission_1);
+        stake::fast_forward_to_unlock(pool_address);
+
+        // Both original stake and operator commissions have received rewards.
+        expected_commission_1 = with_rewards(expected_commission_1);
+        new_balance = with_rewards(new_balance);
+        stake::assert_stake_pool(pool_address, new_balance, expected_commission_1, 0, 0);
+        distribute(staker_address, operator_address);
+        let operator_balance = coin::balance<AptosCoin>(operator_address);
+        let expected_operator_balance = INITIAL_BALANCE + expected_commission_1;
+        assert!(operator_balance == expected_operator_balance, operator_balance);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, 0);
+        assert_no_pending_distributions(staker_address, operator_address);
+
+        // Staker adds more stake.
+        stake::mint(staker, INITIAL_BALANCE);
+        let previous_principal = last_recorded_principal(staker_address, operator_address);
+        add_stake(staker, operator_address, INITIAL_BALANCE);
+        stake::assert_stake_pool(pool_address, new_balance, 0, INITIAL_BALANCE, 0);
+        assert!(last_recorded_principal(staker_address, operator_address) == previous_principal + INITIAL_BALANCE, 0);
+
+        // The newly added stake didn't receive any rewards because it was only added in the new epoch.
+        stake::end_epoch();
+        new_balance = with_rewards(new_balance) + INITIAL_BALANCE;
+
+        // Second round of commission request/withdrawal.
+        let expected_commission_2 = (new_balance - last_recorded_principal(staker_address, operator_address)) / 10;
+        new_balance = new_balance - expected_commission_2;
+        request_commission(operator, staker_address, operator_address);
+        assert_distribution(staker_address, operator_address, operator_address, expected_commission_2);
+        assert!(last_recorded_principal(staker_address, operator_address) == new_balance, 0);
+        stake::fast_forward_to_unlock(pool_address);
+        expected_commission_2 = with_rewards(expected_commission_2);
+        distribute(staker_address, operator_address);
+        operator_balance = coin::balance<AptosCoin>(operator_address);
+        expected_operator_balance = expected_operator_balance + expected_commission_2;
+        assert!(operator_balance == expected_operator_balance, operator_balance);
+        assert_no_pending_distributions(staker_address, operator_address);
+        new_balance = with_rewards(new_balance);
+
+        // New rounds of rewards.
+        stake::fast_forward_to_unlock(pool_address);
+        new_balance = with_rewards(new_balance);
+
+        // Staker withdraws all stake, which should also request commission distribution.
+        let unpaid_commission = (new_balance - last_recorded_principal(staker_address, operator_address)) / 10;
+        unlock_stake(staker, operator_address, new_balance);
+        stake::assert_stake_pool(pool_address, 0, 0, 0, new_balance);
+        assert_distribution(staker_address, operator_address, operator_address, unpaid_commission);
+        let withdrawn_amount = new_balance - unpaid_commission;
+        assert_distribution(staker_address, operator_address, staker_address, withdrawn_amount);
+        assert!(last_recorded_principal(staker_address, operator_address) == 0, 0);
+
+        // End epoch. The stake pool should get kicked out of the validator set as it has 0 remaining active stake.
+        stake::fast_forward_to_unlock(pool_address);
+        // Operator should still earn 10% commission on the rewards on top of the staker's withdrawn_amount.
+        let commission_on_withdrawn_amount = (with_rewards(withdrawn_amount) - withdrawn_amount) / 10;
+        unpaid_commission = with_rewards(unpaid_commission) + commission_on_withdrawn_amount;
+        withdrawn_amount = with_rewards(withdrawn_amount) - commission_on_withdrawn_amount;
+        stake::assert_stake_pool(pool_address, 0, with_rewards(new_balance), 0, 0);
+        assert!(stake::get_validator_state(pool_address) == VALIDATOR_STATUS_INACTIVE, 0);
+
+        // Distribute and verify balances.
+        distribute(staker_address, operator_address);
+        assert_no_pending_distributions(staker_address, operator_address);
+        operator_balance = coin::balance<AptosCoin>(operator_address);
+        assert!(operator_balance == expected_operator_balance + unpaid_commission, operator_balance);
+        let staker_balance = coin::balance<AptosCoin>(staker_address);
+        // Staker receives the extra dust due to rounding error.
+        assert!(staker_balance == withdrawn_amount + 1, staker_balance);
+        assert_no_pending_distributions(staker_address, operator_address);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
+    public entry fun test_operator_cannot_request_same_commission_multiple_times(
+        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, BeneficiaryForOperator {
+        setup_staking_contract(aptos_framework, staker, operator, INITIAL_BALANCE, 10);
+        let staker_address = signer::address_of(staker);
+        let operator_address = signer::address_of(operator);
+        let pool_address = stake_pool_address(staker_address, operator_address);
+
+        // Operator joins the validator set.
+        let (_sk, pk, pop) = stake::generate_identity();
+        stake::join_validator_set_for_test(&pk, &pop, operator, pool_address, true);
+        assert!(stake::get_validator_state(pool_address) == VALIDATOR_STATUS_ACTIVE, 1);
+
+        // Fast forward to generate rewards.
+        stake::end_epoch();
+        let new_balance = with_rewards(INITIAL_BALANCE);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, 0);
+
+        // Operator tries to request commission multiple times. But their distribution shouldn't change.
+        let expected_commission = (new_balance - last_recorded_principal(staker_address, operator_address)) / 10;
+        request_commission(operator, staker_address, operator_address);
+        assert_distribution(staker_address, operator_address, operator_address, expected_commission);
+        request_commission(operator, staker_address, operator_address);
+        assert_distribution(staker_address, operator_address, operator_address, expected_commission);
+        request_commission(operator, staker_address, operator_address);
+        assert_distribution(staker_address, operator_address, operator_address, expected_commission);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
+    public entry fun test_unlock_rewards(
+        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, BeneficiaryForOperator {
+        setup_staking_contract(aptos_framework, staker, operator, INITIAL_BALANCE, 10);
+        let staker_address = signer::address_of(staker);
+        let operator_address = signer::address_of(operator);
+        let pool_address = stake_pool_address(staker_address, operator_address);
+
+        // Operator joins the validator set.
+        let (_sk, pk, pop) = stake::generate_identity();
+        stake::join_validator_set_for_test(&pk, &pop, operator, pool_address, true);
+        assert!(stake::get_validator_state(pool_address) == VALIDATOR_STATUS_ACTIVE, 1);
+
+        // Fast forward to generate rewards.
+        stake::end_epoch();
+        let new_balance = with_rewards(INITIAL_BALANCE);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, 0);
+
+        // Staker withdraws all accumulated rewards, which should pay commission first.
+        unlock_rewards(staker, operator_address);
+        let accumulated_rewards = new_balance - INITIAL_BALANCE;
+        let expected_commission = accumulated_rewards / 10;
+        let staker_rewards = accumulated_rewards - expected_commission;
+        assert_distribution(staker_address, operator_address, staker_address, staker_rewards);
+        assert_distribution(staker_address, operator_address, operator_address, expected_commission);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
+    #[expected_failure(abort_code = 0x80006, location = Self)]
+    public entry fun test_staker_cannot_create_same_staking_contract_multiple_times(
+        aptos_framework: &signer,
+        staker: &signer,
+        operator: &signer,
+    ) acquires Store {
+        setup_staking_contract(aptos_framework, staker, operator, INITIAL_BALANCE, 10);
+        let operator_address = signer::address_of(operator);
+        stake::mint(staker, INITIAL_BALANCE);
+        create_staking_contract(staker, operator_address, operator_address, INITIAL_BALANCE, 10, vector::empty<u8>());
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
+    #[expected_failure(abort_code = 0x10002, location = Self)]
+    public entry fun test_staker_cannot_create_staking_contract_with_invalid_commission(
+        aptos_framework: &signer,
+        staker: &signer,
+        operator: &signer,
+    ) acquires Store {
+        setup_staking_contract(aptos_framework, staker, operator, INITIAL_BALANCE, 101);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
+    #[expected_failure(abort_code = 0x10001, location = Self)]
+    public entry fun test_staker_cannot_create_staking_contract_with_less_than_min_stake_required(
+        aptos_framework: &signer,
+        staker: &signer,
+        operator: &signer,
+    ) acquires Store {
+        setup_staking_contract(aptos_framework, staker, operator, 50, 100);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
+    public entry fun test_update_voter(
+        aptos_framework: &signer,
+        staker: &signer,
+        operator: &signer,
+    ) acquires Store {
+        setup_staking_contract(aptos_framework, staker, operator, INITIAL_BALANCE, 10);
+        let staker_address = signer::address_of(staker);
+        let operator_address = signer::address_of(operator);
+
+        // Voter is initially set to operator but then updated to be staker.
+        let pool_address = stake_pool_address(staker_address, operator_address);
+        assert!(stake::get_delegated_voter(pool_address) == operator_address, 0);
+        update_voter(staker, operator_address, staker_address);
+        assert!(stake::get_delegated_voter(pool_address) == staker_address, 1);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
+    public entry fun test_reset_lockup(
+        aptos_framework: &signer,
+        staker: &signer,
+        operator: &signer,
+    ) acquires Store {
+        setup_staking_contract(aptos_framework, staker, operator, INITIAL_BALANCE, 10);
+        let staker_address = signer::address_of(staker);
+        let operator_address = signer::address_of(operator);
+        let pool_address = stake_pool_address(staker_address, operator_address);
+
+        let origin_lockup_expiration = stake::get_lockup_secs(pool_address);
+        reset_lockup(staker, operator_address);
+        assert!(origin_lockup_expiration < stake::get_lockup_secs(pool_address), 0);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator_1 = @0x234, operator_2 = @0x345)]
+    public entry fun test_staker_can_switch_operator(
+        aptos_framework: &signer,
+        staker: &signer,
+        operator_1: &signer,
+        operator_2: &signer,
+    ) acquires Store, BeneficiaryForOperator {
+        setup_staking_contract(aptos_framework, staker, operator_1, INITIAL_BALANCE, 10);
+        account::create_account_for_test(signer::address_of(operator_2));
+        stake::mint(operator_2, INITIAL_BALANCE);
+        let staker_address = signer::address_of(staker);
+        let operator_1_address = signer::address_of(operator_1);
+        let operator_2_address = signer::address_of(operator_2);
+
+        // Join validator set and earn some rewards.
+        let pool_address = stake_pool_address(staker_address, operator_1_address);
+        let (_sk, pk, pop) = stake::generate_identity();
+        stake::join_validator_set_for_test(&pk, &pop, operator_1, pool_address, true);
+        stake::end_epoch();
+        assert!(stake::get_validator_state(pool_address) == VALIDATOR_STATUS_ACTIVE, 0);
+
+        // Switch operators.
+        switch_operator(staker, operator_1_address, operator_2_address, 20);
+        // The staking_contract is now associated with operator 2 but there should be a pending distribution of unpaid
+        // commission to operator 1.
+        let new_balance = with_rewards(INITIAL_BALANCE);
+        let commission_for_operator_1 = (new_balance - INITIAL_BALANCE) / 10;
+        assert_distribution(staker_address, operator_2_address, operator_1_address, commission_for_operator_1);
+        // Unpaid commission should be unlocked from the stake pool.
+        new_balance = new_balance - commission_for_operator_1;
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, commission_for_operator_1);
+        assert!(last_recorded_principal(staker_address, operator_2_address) == new_balance, 0);
+
+        // The stake pool's validator should not have left the validator set.
+        assert!(stake_pool_address(staker_address, operator_2_address) == pool_address, 1);
+        assert!(stake::get_validator_state(pool_address) == VALIDATOR_STATUS_ACTIVE, 2);
+
+        // End epoch to get more rewards.
+        stake::fast_forward_to_unlock(pool_address);
+        new_balance = with_rewards(new_balance);
+        // Rewards on the commission being paid to operator_1 should still be charged commission that will go to
+        // operator_2;
+        let commission_on_operator_1_distribution =
+            (with_rewards(commission_for_operator_1) - commission_for_operator_1) / 5;
+        commission_for_operator_1 = with_rewards(commission_for_operator_1) - commission_on_operator_1_distribution;
+
+        // Verify that when commissions are withdrawn, previous pending distribution to operator 1 also happens.
+        // Then new commission of 20% is paid to operator 2.
+        let commission_for_operator_2 =
+            (new_balance - last_recorded_principal(staker_address, operator_2_address)) / 5;
+        new_balance = new_balance - commission_for_operator_2;
+        request_commission(operator_2, staker_address, operator_2_address);
+        assert_distribution(staker_address, operator_2_address, operator_2_address, commission_for_operator_2);
+        let operator_1_balance = coin::balance<AptosCoin>(operator_1_address);
+        assert!(operator_1_balance == INITIAL_BALANCE + commission_for_operator_1, operator_1_balance);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, commission_for_operator_2);
+        assert!(last_recorded_principal(staker_address, operator_2_address) == new_balance, 0);
+        stake::fast_forward_to_unlock(pool_address);
+
+        // Operator 2's commission is distributed.
+        distribute(staker_address, operator_2_address);
+        let operator_2_balance = coin::balance<AptosCoin>(operator_2_address);
+        new_balance = with_rewards(new_balance);
+        commission_for_operator_2 = with_rewards(commission_for_operator_2);
+        assert!(
+            operator_2_balance == INITIAL_BALANCE + commission_for_operator_2 + commission_on_operator_1_distribution,
+            operator_2_balance,
+        );
+        stake::assert_stake_pool(
+            pool_address,
+            new_balance,
+            0,
+            0,
+            0,
+        );
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator_1 = @0x234, operator_2 = @0x345)]
+    public entry fun test_staker_can_switch_operator_with_same_commission(
+        aptos_framework: &signer,
+        staker: &signer,
+        operator_1: &signer,
+        operator_2: &signer,
+    ) acquires Store, BeneficiaryForOperator {
+        setup_staking_contract(aptos_framework, staker, operator_1, INITIAL_BALANCE, 10);
+        let staker_address = signer::address_of(staker);
+        let operator_1_address = signer::address_of(operator_1);
+        let operator_2_address = signer::address_of(operator_2);
+
+        // Switch operators.
+        switch_operator_with_same_commission(staker, operator_1_address, operator_2_address);
+        // The staking_contract should now be associated with operator 2 but with same commission rate.
+        assert!(staking_contract_exists(staker_address, operator_2_address), 0);
+        assert!(!staking_contract_exists(staker_address, operator_1_address), 1);
+        assert!(commission_percentage(staker_address, operator_2_address) == 10, 2);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator1 = @0x234, beneficiary = @0x345, operator2 = @0x456)]
+    public entry fun test_operator_can_set_beneficiary(
+        aptos_framework: &signer,
+        staker: &signer,
+        operator1: &signer,
+        beneficiary: &signer,
+        operator2: &signer,
+    ) acquires Store, BeneficiaryForOperator {
+        setup_staking_contract(aptos_framework, staker, operator1, INITIAL_BALANCE, 10);
+        let staker_address = signer::address_of(staker);
+        let operator1_address = signer::address_of(operator1);
+        let operator2_address = signer::address_of(operator2);
+        let beneficiary_address = signer::address_of(beneficiary);
+
+        // account::create_account_for_test(beneficiary_address);
+        aptos_framework::aptos_account::create_account(beneficiary_address);
+        assert_staking_contract_exists(staker_address, operator1_address);
+        assert_staking_contract(staker_address, operator1_address, INITIAL_BALANCE, 10);
+
+        // Verify that the stake pool has been set up properly.
+        let pool_address = stake_pool_address(staker_address, operator1_address);
+        stake::assert_stake_pool(pool_address, INITIAL_BALANCE, 0, 0, 0);
+        assert!(last_recorded_principal(staker_address, operator1_address) == INITIAL_BALANCE, 0);
+        assert!(stake::get_operator(pool_address) == operator1_address, 0);
+        assert!(beneficiary_for_operator(operator1_address) == operator1_address, 0);
+
+        // Operator joins the validator set.
+        let (_sk, pk, pop) = stake::generate_identity();
+        stake::join_validator_set_for_test(&pk, &pop, operator1, pool_address, true);
+        assert!(stake::get_validator_state(pool_address) == VALIDATOR_STATUS_ACTIVE, 1);
+
+        // Set beneficiary.
+        set_beneficiary_for_operator(operator1, beneficiary_address);
+        assert!(beneficiary_for_operator(operator1_address) == beneficiary_address, 0);
+
+        // Fast forward to generate rewards.
+        stake::end_epoch();
+        let new_balance = with_rewards(INITIAL_BALANCE);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, 0);
+
+        // Operator claims 10% of rewards so far as commissions.
+        let expected_commission_1 = (new_balance - last_recorded_principal(staker_address, operator1_address)) / 10;
+        new_balance = new_balance - expected_commission_1;
+        request_commission(operator1, staker_address, operator1_address);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, expected_commission_1);
+        assert!(last_recorded_principal(staker_address, operator1_address) == new_balance, 0);
+        assert_distribution(staker_address, operator1_address, operator1_address, expected_commission_1);
+        stake::fast_forward_to_unlock(pool_address);
+
+        // Both original stake and operator commissions have received rewards.
+        expected_commission_1 = with_rewards(expected_commission_1);
+        new_balance = with_rewards(new_balance);
+        stake::assert_stake_pool(pool_address, new_balance, expected_commission_1, 0, 0);
+        distribute(staker_address, operator1_address);
+        let operator_balance = coin::balance<AptosCoin>(operator1_address);
+        let beneficiary_balance = coin::balance<AptosCoin>(beneficiary_address);
+        let expected_operator_balance = INITIAL_BALANCE;
+        let expected_beneficiary_balance = expected_commission_1;
+        assert!(operator_balance == expected_operator_balance, operator_balance);
+        assert!(beneficiary_balance == expected_beneficiary_balance, beneficiary_balance);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, 0);
+        assert_no_pending_distributions(staker_address, operator1_address);
+
+        // switch operator to operator2. The rewards should go to operator2 not to the beneficiay of operator1.
+        let old_beneficiay_balance = beneficiary_balance;
+        switch_operator(staker, operator1_address, operator2_address, 10);
+
+        stake::end_epoch();
+        let (_, accumulated_rewards, _) = staking_contract_amounts(staker_address, operator2_address);
+
+        let expected_commission = accumulated_rewards / 10;
+
+        // Request commission.
+        request_commission(operator2, staker_address, operator2_address);
+        // Unlocks the commission.
+        stake::fast_forward_to_unlock(pool_address);
+        expected_commission = with_rewards(expected_commission);
+
+        // Distribute the commission to the operator.
+        distribute(staker_address, operator2_address);
+
+        // Assert that the rewards go to operator2, and the balance of the operator1's beneficiay remains the same.
+        assert!(coin::balance<AptosCoin>(operator2_address) >= expected_commission, 1);
+        assert!(coin::balance<AptosCoin>(beneficiary_address) == old_beneficiay_balance, 1);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
+    public entry fun test_staker_can_withdraw_partial_stake(
+        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, BeneficiaryForOperator {
+        let initial_balance = INITIAL_BALANCE * 2;
+        setup_staking_contract(aptos_framework, staker, operator, initial_balance, 10);
+        let staker_address = signer::address_of(staker);
+        let operator_address = signer::address_of(operator);
+        let pool_address = stake_pool_address(staker_address, operator_address);
+
+        // Operator joins the validator set so rewards are generated.
+        let (_sk, pk, pop) = stake::generate_identity();
+        stake::join_validator_set_for_test(&pk, &pop, operator, pool_address, true);
+        assert!(stake::get_validator_state(pool_address) == VALIDATOR_STATUS_ACTIVE, 1);
+
+        // Fast forward to generate rewards.
+        stake::end_epoch();
+        let new_balance = with_rewards(initial_balance);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, 0);
+
+        // Staker withdraws 1/4 of the stake, which should also request commission distribution.
+        let withdrawn_stake = new_balance / 4;
+        let unpaid_commission = (new_balance - initial_balance) / 10;
+        let new_balance = new_balance - withdrawn_stake - unpaid_commission;
+        unlock_stake(staker, operator_address, withdrawn_stake);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, withdrawn_stake + unpaid_commission);
+        assert_distribution(staker_address, operator_address, operator_address, unpaid_commission);
+        assert_distribution(staker_address, operator_address, staker_address, withdrawn_stake);
+        assert!(last_recorded_principal(staker_address, operator_address) == new_balance, 0);
+
+        // The validator is still in the active set as its remaining stake is still above min required.
+        stake::fast_forward_to_unlock(pool_address);
+        new_balance = with_rewards(new_balance);
+        unpaid_commission = with_rewards(unpaid_commission);
+        // Commission should still be charged on the rewards on top of withdrawn_stake.
+        // So the operator should receive 10% of the rewards on top of withdrawn_stake.
+        let commission_on_withdrawn_stake = (with_rewards(withdrawn_stake) - withdrawn_stake) / 10;
+        unpaid_commission = unpaid_commission + commission_on_withdrawn_stake;
+        withdrawn_stake = with_rewards(withdrawn_stake) - commission_on_withdrawn_stake;
+        stake::assert_stake_pool(pool_address, new_balance, withdrawn_stake + unpaid_commission, 0, 0);
+        assert!(stake::get_validator_state(pool_address) == VALIDATOR_STATUS_ACTIVE, 0);
+
+        // Distribute and verify balances.
+        distribute(staker_address, operator_address);
+        assert_no_pending_distributions(staker_address, operator_address);
+        let operator_balance = coin::balance<AptosCoin>(operator_address);
+        assert!(operator_balance == initial_balance + unpaid_commission, operator_balance);
+        let staker_balance = coin::balance<AptosCoin>(staker_address);
+        assert!(staker_balance == withdrawn_stake, staker_balance);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
+    public entry fun test_staker_can_withdraw_partial_stake_if_operator_never_joined_validator_set(
+        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, BeneficiaryForOperator {
+        let initial_balance = INITIAL_BALANCE * 2;
+        setup_staking_contract(aptos_framework, staker, operator, initial_balance, 10);
+        let staker_address = signer::address_of(staker);
+        let operator_address = signer::address_of(operator);
+        let pool_address = stake_pool_address(staker_address, operator_address);
+
+        // Epoch ended, but since validator never joined the set, no rewards were minted.
+        stake::end_epoch();
+        stake::assert_stake_pool(pool_address, initial_balance, 0, 0, 0);
+
+        // Staker withdraws 1/4 of the stake, which doesn't create any commission distribution as there's no rewards.
+        let withdrawn_stake = initial_balance / 4;
+        let new_balance = initial_balance - withdrawn_stake;
+        unlock_stake(staker, operator_address, withdrawn_stake);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, withdrawn_stake);
+        assert_distribution(staker_address, operator_address, operator_address, 0);
+        assert_distribution(staker_address, operator_address, staker_address, withdrawn_stake);
+        assert!(last_recorded_principal(staker_address, operator_address) == new_balance, 0);
+
+        // Distribute and verify balances.
+        distribute(staker_address, operator_address);
+        assert_no_pending_distributions(staker_address, operator_address);
+        // Operator's balance shouldn't change as there are no rewards.
+        let operator_balance = coin::balance<AptosCoin>(operator_address);
+        assert!(operator_balance == initial_balance, operator_balance);
+        // Staker receives back the withdrawn amount (no rewards).
+        let staker_balance = coin::balance<AptosCoin>(staker_address);
+        assert!(staker_balance == withdrawn_stake, staker_balance);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
+    public entry fun test_multiple_distributions_added_before_distribute(
+        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, BeneficiaryForOperator {
+        let initial_balance = INITIAL_BALANCE * 2;
+        setup_staking_contract(aptos_framework, staker, operator, initial_balance, 10);
+        let staker_address = signer::address_of(staker);
+        let operator_address = signer::address_of(operator);
+        let pool_address = stake_pool_address(staker_address, operator_address);
+
+        // Operator joins the validator set so rewards are generated.
+        let (_sk, pk, pop) = stake::generate_identity();
+        stake::join_validator_set_for_test(&pk, &pop, operator, pool_address, true);
+        assert!(stake::get_validator_state(pool_address) == VALIDATOR_STATUS_ACTIVE, 1);
+
+        // Fast forward to generate rewards.
+        stake::end_epoch();
+        let new_balance = with_rewards(initial_balance);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, 0);
+
+        // Staker withdraws 1/4 of the stake, which should also request commission distribution.
+        let withdrawn_stake = new_balance / 4;
+        let unpaid_commission = (new_balance - initial_balance) / 10;
+        let new_balance = new_balance - withdrawn_stake - unpaid_commission;
+        unlock_stake(staker, operator_address, withdrawn_stake);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, withdrawn_stake + unpaid_commission);
+        assert_distribution(staker_address, operator_address, operator_address, unpaid_commission);
+        assert_distribution(staker_address, operator_address, staker_address, withdrawn_stake);
+        assert!(last_recorded_principal(staker_address, operator_address) == new_balance, 0);
+
+        // End epoch to generate some rewards. Staker withdraws another 1/4 of the stake.
+        // Commission should be charged on the rewards earned on the previous 1/4 stake withdrawal.
+        stake::end_epoch();
+        let commission_on_withdrawn_stake = (with_rewards(withdrawn_stake) - withdrawn_stake) / 10;
+        let commission_on_new_balance = (with_rewards(new_balance) - new_balance) / 10;
+        unpaid_commission = with_rewards(unpaid_commission) + commission_on_withdrawn_stake + commission_on_new_balance;
+        new_balance = with_rewards(new_balance) - commission_on_new_balance;
+        let new_withdrawn_stake = new_balance / 4;
+        unlock_stake(staker, operator_address, new_withdrawn_stake);
+        new_balance = new_balance - new_withdrawn_stake;
+        withdrawn_stake = with_rewards(withdrawn_stake) - commission_on_withdrawn_stake + new_withdrawn_stake;
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, withdrawn_stake + unpaid_commission);
+        // There's some small rounding error here.
+        assert_distribution(staker_address, operator_address, operator_address, unpaid_commission - 1);
+        assert_distribution(staker_address, operator_address, staker_address, withdrawn_stake);
+        assert!(last_recorded_principal(staker_address, operator_address) == new_balance, 0);
+    }
+
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
+    public entry fun test_update_commission(
+        aptos_framework: &signer,
+        staker: &signer,
+        operator: &signer
+    ) acquires Store, BeneficiaryForOperator, StakingGroupUpdateCommissionEvent {
+        let initial_balance = INITIAL_BALANCE * 2;
+        setup_staking_contract(aptos_framework, staker, operator, initial_balance, 10);
+        let staker_address = signer::address_of(staker);
+        let operator_address = signer::address_of(operator);
+        let pool_address = stake_pool_address(staker_address, operator_address);
+
+        // Operator joins the validator set so rewards are generated.
+        let (_sk, pk, pop) = stake::generate_identity();
+        stake::join_validator_set_for_test(&pk, &pop, operator, pool_address, true);
+        assert!(stake::get_validator_state(pool_address) == VALIDATOR_STATUS_ACTIVE, 1);
+
+        // Fast forward to generate rewards.
+        stake::end_epoch();
+        let balance_1epoch = with_rewards(initial_balance);
+        let unpaid_commission = (balance_1epoch - initial_balance) / 10;
+        stake::assert_stake_pool(pool_address, balance_1epoch, 0, 0, 0);
+
+        update_commision(staker, operator_address, 5);
+        stake::end_epoch();
+        let balance_2epoch = with_rewards(balance_1epoch - unpaid_commission);
+        stake::assert_stake_pool(pool_address, balance_2epoch, 0, 0, with_rewards(unpaid_commission));
+    }
+
+    #[test(
+        staker = @0xe256f4f4e2986cada739e339895cf5585082ff247464cab8ec56eea726bd2263,
+        operator = @0x9f0a211d218b082987408f1e393afe1ba0c202c6d280f081399788d3360c7f09
+    )]
+    public entry fun test_get_expected_stake_pool_address(staker: address, operator: address) {
+        let pool_address = get_expected_stake_pool_address(staker, operator, vector[0x42, 0x42]);
+        assert!(pool_address == @0x9d9648031ada367c26f7878eb0b0406ae6a969b1a43090269e5cdfabe1b48f0f, 0);
+    }
+
+    #[test_only]
+    public fun assert_staking_contract(
+        staker: address, operator: address, principal: u64, commission_percentage: u64) acquires Store {
+        let staking_contract = simple_map::borrow(&borrow_global<Store>(staker).staking_contracts, &operator);
+        assert!(staking_contract.principal == principal, staking_contract.principal);
+        assert!(
+            staking_contract.commission_percentage == commission_percentage,
+            staking_contract.commission_percentage
+        );
+    }
+
+    #[test_only]
+    public fun assert_no_pending_distributions(staker: address, operator: address) acquires Store {
+        let staking_contract = simple_map::borrow(&borrow_global<Store>(staker).staking_contracts, &operator);
+        let distribution_pool = &staking_contract.distribution_pool;
+        let shareholders_count = pool_u64::shareholders_count(distribution_pool);
+        assert!(shareholders_count == 0, shareholders_count);
+        let total_coins_remaining = pool_u64::total_coins(distribution_pool);
+        assert!(total_coins_remaining == 0, total_coins_remaining);
+    }
+
+    #[test_only]
+    public fun assert_distribution(
+        staker: address, operator: address, recipient: address, coins_amount: u64) acquires Store {
+        let staking_contract = simple_map::borrow(&borrow_global<Store>(staker).staking_contracts, &operator);
+        let distribution_balance = pool_u64::balance(&staking_contract.distribution_pool, recipient);
+        assert!(distribution_balance == coins_amount, distribution_balance);
+    }
+}