Fix cont.

Author: Dinonard

pallets/dapp-staking-v3/src/test/testing_utils.rs

@@ -627,17 +627,6 @@ pub(crate) fn assert_unstake(
         .expect("Entry must exist since 'unstake' is being called.");
     let pre_era_info = pre_snapshot.current_era_info;
 
-    // Calculate expected unstake amounts from the voting and build&earn subperiods.
-    let (expected_voting_unstake, expected_build_and_earn_unstake) =
-        if pre_staker_info.staked_amount(Subperiod::BuildAndEarn) >= amount {
-            (Balance::zero(), amount)
-        } else {
-            (
-                amount - pre_staker_info.staked_amount(Subperiod::BuildAndEarn),
-                pre_staker_info.staked_amount(Subperiod::BuildAndEarn),
-            )
-        };
-
     let unstake_period = pre_snapshot.active_protocol_state.period_number();
     let unstake_subperiod = pre_snapshot.active_protocol_state.subperiod();
 
@@ -652,6 +641,17 @@ pub(crate) fn assert_unstake(
         amount
     };
 
+    // Calculate expected unstake amounts from the voting and build&earn subperiods.
+    let (expected_voting_unstake, expected_build_and_earn_unstake) =
+        if pre_staker_info.staked_amount(Subperiod::BuildAndEarn) >= amount {
+            (Balance::zero(), amount)
+        } else {
+            (
+                amount - pre_staker_info.staked_amount(Subperiod::BuildAndEarn),
+                pre_staker_info.staked_amount(Subperiod::BuildAndEarn),
+            )
+        };
+
     // Unstake from smart contract & verify event
     assert_ok!(DappStaking::unstake(
         RuntimeOrigin::signed(account),
@@ -747,40 +747,18 @@ pub(crate) fn assert_unstake(
             .saturating_sub(amount),
         "Staked amount must decreased by the 'amount'"
     );
+
     assert_eq!(
-        post_contract_stake.staked.voting,
-        pre_contract_stake.staked.voting.saturating_sub(expected_voting_unstake),
+        post_contract_stake.staked_amount(unstake_period, Subperiod::Voting),
+        pre_contract_stake.staked_amount(unstake_period, Subperiod::Voting) - expected_voting_unstake,
         "Voting stake amount must be decreased exactly by how much staker's voting amount was reduced."
     );
-
     assert_eq!(
-        post_contract_stake.staked.build_and_earn,
-        pre_contract_stake
-            .staked
-            .build_and_earn
-            .saturating_sub(expected_build_and_earn_unstake),
+        post_contract_stake.staked_amount(unstake_period, Subperiod::BuildAndEarn),
+        pre_contract_stake.staked_amount(unstake_period, Subperiod::BuildAndEarn)
+            - expected_build_and_earn_unstake,
         "B&E stake amount must be decreased exactly by how much staker's B&E amount was reduced."
     );
-    match (
-        pre_contract_stake.staked_future,
-        post_contract_stake.staked_future,
-    ) {
-        (Some(pre_future), Some(post_future)) => {
-            assert_eq!(
-                post_future.voting,
-                pre_future.voting.saturating_sub(expected_voting_unstake),
-                "Voting stake amount must be decreased exactly by how much staker's voting amount was reduced."
-            );
-            assert_eq!(
-                post_future.build_and_earn,
-                pre_future
-                    .build_and_earn
-                    .saturating_sub(expected_build_and_earn_unstake),
-                "B&E stake amount must be decreased exactly by how much staker's B&E amount was reduced."
-            );
-        }
-        (_, _) => {}
-    }
 
     // 4. verify era info
     // =========================

pallets/dapp-staking-v3/src/test/tests_types.rs

@@ -2430,32 +2430,42 @@ fn contract_stake_amount_stake_is_ok() {
     assert!(contract_stake.staked_future.is_some());
 }
 
-// TODO: improve this a bit
 #[test]
 fn contract_stake_amount_unstake_is_ok() {
     let mut contract_stake = ContractStakeAmount::default();
 
     // Prep action - create a stake entry
     let era_1 = 2;
+    let era_2 = era_1 + 1;
     let period = 3;
     let period_info = PeriodInfo {
         number: period,
         subperiod: Subperiod::Voting,
         next_subperiod_start_era: 20,
     };
-    let stake_amount = 100;
-    contract_stake.stake(stake_amount, period_info, era_1);
+    let vp_stake_amount = 47;
+    let bep_stake_amount = 53;
+    contract_stake.stake(vp_stake_amount, period_info, era_1);
+    contract_stake.stake(
+        bep_stake_amount,
+        PeriodInfo {
+            subperiod: Subperiod::BuildAndEarn,
+            ..period_info
+        },
+        era_1,
+    );
+    let total_stake_amount = vp_stake_amount + bep_stake_amount;
 
-    // 1st scenario - unstake in the same era
+    // 1st scenario - unstake in the same era, from `Voting` subperiod
     let amount_1 = 5;
     contract_stake.unstake(amount_1, 0, period_info, era_1);
     assert_eq!(
         contract_stake.total_staked_amount(period),
-        stake_amount - amount_1
+        total_stake_amount - amount_1
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        stake_amount - amount_1
+        vp_stake_amount - amount_1
     );
     assert!(contract_stake.staked.is_empty());
     assert!(contract_stake.staked_future.is_some());
@@ -2466,16 +2476,20 @@ fn contract_stake_amount_unstake_is_ok() {
         subperiod: Subperiod::BuildAndEarn,
         next_subperiod_start_era: 40,
     };
-    let era_2 = era_1 + 1;
 
-    contract_stake.unstake(0, amount_1, period_info, era_2);
+    contract_stake.unstake(amount_1, 0, period_info, era_2);
     assert_eq!(
         contract_stake.total_staked_amount(period),
-        stake_amount - amount_1 * 2
+        total_stake_amount - amount_1 * 2
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        stake_amount - amount_1 * 2
+        vp_stake_amount - amount_1 * 2
+    );
+    assert_eq!(
+        contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
+        bep_stake_amount,
+        "Must remain unchanged."
     );
     assert!(
         !contract_stake.staked.is_empty(),
@@ -2486,17 +2500,36 @@ fn contract_stake_amount_unstake_is_ok() {
         "future entry should be cleaned up since it refers to the current era"
     );
 
-    // 3rd scenario - bump up unstake eras by more than 1, entries should be aligned to the current era
+    // 3rd scenario - same as previous scenario, but for `BuildAndEarn` subperiod
+    contract_stake.unstake(0, amount_1, period_info, era_2);
+    assert_eq!(
+        contract_stake.total_staked_amount(period),
+        total_stake_amount - amount_1 * 3
+    );
+    assert_eq!(
+        contract_stake.staked_amount(period, Subperiod::Voting),
+        vp_stake_amount - amount_1 * 2
+    );
+    assert_eq!(
+        contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
+        bep_stake_amount - amount_1
+    );
+
+    // 4th scenario - bump up unstake eras by more than 1, entries should be aligned to the current era
     let era_3 = era_2 + 3;
     let amount_2 = 7;
     contract_stake.unstake(0, amount_2, period_info, era_3);
     assert_eq!(
         contract_stake.total_staked_amount(period),
-        stake_amount - amount_1 * 2 - amount_2
+        total_stake_amount - amount_1 * 3 - amount_2
     );
     assert_eq!(
         contract_stake.staked_amount(period, Subperiod::Voting),
-        stake_amount - amount_1 * 2 - amount_2
+        vp_stake_amount - amount_1 * 2
+    );
+    assert_eq!(
+        contract_stake.staked_amount(period, Subperiod::BuildAndEarn),
+        bep_stake_amount - amount_1 - amount_2
     );
     assert_eq!(
         contract_stake.staked.era, era_3,
@@ -2507,8 +2540,8 @@ fn contract_stake_amount_unstake_is_ok() {
         "future entry should remain 'None'"
     );
 
-    // 4th scenario - do a full unstake with existing future entry, expect a cleanup
-    contract_stake.stake(stake_amount, period_info, era_3);
+    // 5th scenario - do a full unstake with existing future entry, expect a cleanup
+    contract_stake.stake(total_stake_amount, period_info, era_3);
     contract_stake.unstake(
         contract_stake.staked_amount(period, Subperiod::Voting),
         contract_stake.staked_amount(period, Subperiod::BuildAndEarn),

pallets/dapp-staking-v3/src/types.rs

@@ -882,6 +882,13 @@ impl StakeAmount {
             }
         }
     }
+
+    /// Saturating reduce the specified `amount` from the appropriate subperiod.
+    /// If the amount is larger than the subperiod amount, the subperiod amount is set to zero.
+    pub fn subtract_subperiods(&mut self, voting_amount: Balance, bep_amount: Balance) {
+        self.voting.saturating_reduce(voting_amount);
+        self.build_and_earn.saturating_reduce(bep_amount);
+    }
 }
 
 /// Info about an era, including the rewards, how much is locked, unlocking, etc.
@@ -1198,7 +1205,7 @@ impl ContractStakeAmount {
         }
     }
 
-    /// Unstake the specified `amount` from the contract, for the specified `subperiod` and `era`.
+    /// Unstake the specified `Voting` and `Build&Earn` subperiod amounts from the contract, for the specified `subperiod` and `era`.
     pub fn unstake(
         &mut self,
         voting_amount: Balance,
@@ -1225,11 +1232,9 @@ impl ContractStakeAmount {
         }
 
         // Subtract both amounts
-        self.staked.voting.saturating_reduce(voting_amount);
-        self.staked.build_and_earn.saturating_reduce(bep_amount);
+        self.staked.subtract_subperiods(voting_amount, bep_amount);
         if let Some(stake_amount) = self.staked_future.as_mut() {
-            stake_amount.voting.saturating_reduce(voting_amount);
-            stake_amount.build_and_earn.saturating_reduce(bep_amount);
+            stake_amount.subtract_subperiods(voting_amount, bep_amount);
         }
 
         // Convenience cleanup