From 2d2776a7fbe1dfab55a11b52e5ec5726be0a0c71 Mon Sep 17 00:00:00 2001
From: Dino Pacandi <dino.pacandi@gmail.com>
Date: Fri, 21 Jun 2024 15:14:01 +0200
Subject: [PATCH] Fix incorrect check for locked/frozen amount

---
 pallets/dapp-staking-v3/src/lib.rs               |  2 +-
 .../dapp-staking-v3/src/test/testing_utils.rs    | 13 +++++++++----
 pallets/dapp-staking-v3/src/test/tests.rs        | 16 ++++++++++++++++
 3 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/pallets/dapp-staking-v3/src/lib.rs b/pallets/dapp-staking-v3/src/lib.rs
index efd59cf16..c1cce8f8a 100644
--- a/pallets/dapp-staking-v3/src/lib.rs
+++ b/pallets/dapp-staking-v3/src/lib.rs
@@ -838,7 +838,7 @@ pub mod pallet {
 
             // Calculate & check amount available for locking
             let available_balance =
-                T::Currency::total_balance(&account).saturating_sub(ledger.active_locked_amount());
+                T::Currency::total_balance(&account).saturating_sub(ledger.total_locked_amount());
             let amount_to_lock = available_balance.min(amount);
             ensure!(!amount_to_lock.is_zero(), Error::<T>::ZeroAmount);
 
diff --git a/pallets/dapp-staking-v3/src/test/testing_utils.rs b/pallets/dapp-staking-v3/src/test/testing_utils.rs
index 9ecab7957..018593e6b 100644
--- a/pallets/dapp-staking-v3/src/test/testing_utils.rs
+++ b/pallets/dapp-staking-v3/src/test/testing_utils.rs
@@ -86,9 +86,9 @@ impl MemorySnapshot {
     /// Returns locked balance in dApp staking for the specified account.
     /// In case no balance is locked, returns zero.
     pub fn locked_balance(&self, account: &AccountId) -> Balance {
-        self.ledger
-            .get(&account)
-            .map_or(Balance::zero(), |ledger| ledger.active_locked_amount())
+        self.ledger.get(&account).map_or(Balance::zero(), |ledger| {
+            ledger.locked + ledger.unlocking.iter().fold(0, |acc, x| acc + x.amount)
+        })
     }
 }
 
@@ -240,10 +240,15 @@ pub(crate) fn assert_lock(account: AccountId, amount: Balance) {
         "Total locked balance should be increased by the amount locked."
     );
 
+    let post_frozen_balance = Balances::balance_frozen(&FreezeReason::DAppStaking.into(), &account);
     assert_eq!(
         init_frozen_balance + expected_lock_amount,
-        Balances::balance_frozen(&FreezeReason::DAppStaking.into(), &account)
+        post_frozen_balance
     );
+    assert!(
+        Balances::total_balance(&account) >= post_frozen_balance,
+        "Total balance should never be less than frozen balance."
+    )
 }
 
 /// Start the unlocking process for locked funds and assert success.
diff --git a/pallets/dapp-staking-v3/src/test/tests.rs b/pallets/dapp-staking-v3/src/test/tests.rs
index f1bc1cf0d..ad7a0681b 100644
--- a/pallets/dapp-staking-v3/src/test/tests.rs
+++ b/pallets/dapp-staking-v3/src/test/tests.rs
@@ -3286,3 +3286,19 @@ fn unstake_correctly_reduces_future_contract_stake() {
         assert_unstake(staker_1, &smart_contract, amount_2 + 3);
     })
 }
+
+#[test]
+fn lock_correctly_considers_unlocking_amount() {
+    ExtBuilder::build().execute_with(|| {
+        // Lock the entire amount & immediately start the unlocking process
+        let (staker, unlock_amount) = (1, 13);
+        let total_balance = Balances::total_balance(&staker);
+        assert_lock(staker, total_balance);
+        assert_unlock(staker, unlock_amount);
+
+        assert_noop!(
+            DappStaking::lock(RuntimeOrigin::signed(staker), 1),
+            Error::<Test>::ZeroAmount
+        );
+    })
+}