diff --git a/substrate/frame/staking-rewards/src/lib.rs b/substrate/frame/staking-rewards/src/lib.rs
index b74d4b50eb4635f3126f42c7935252eb3794f174..1883b1dffba04535c8ad93cf55e6c8c47f906d80 100644
--- a/substrate/frame/staking-rewards/src/lib.rs
+++ b/substrate/frame/staking-rewards/src/lib.rs
@@ -238,6 +238,8 @@ pub mod pallet {
#[pallet::error]
pub enum Error<T> {
+ /// The staker does not have enough tokens to perform the operation.
+ NotEnoughTokens,
/// An operation was attempted on a non-existent pool.
NonExistentPool,
/// An operation was attempted using a non-existent asset.
@@ -357,16 +359,21 @@ pub mod pallet {
// Always start by updating the pool rewards.
Self::update_pool_rewards(&pool_id, &caller)?;
+ // Check the staker has enough staked tokens.
+ let mut staker = PoolStakers::<T>::get(pool_id, &caller).unwrap_or_default();
+ ensure!(staker.amount >= amount, Error::<T>::NotEnoughTokens);
+
// Unfreeze staker assets.
// TODO: (blocked https://github.com/paritytech/polkadot-sdk/issues/3342)
// Update Pools.
let mut pool = Pools::<T>::get(pool_id).ok_or(Error::<T>::NonExistentPool)?;
pool.total_tokens_staked.saturating_reduce(amount);
+ Pools::<T>::insert(pool_id, pool);
// Update PoolStakers.
- let mut staker = PoolStakers::<T>::get(pool_id, &caller).unwrap_or_default();
staker.amount.saturating_reduce(amount);
+ PoolStakers::<T>::insert(pool_id, &caller, staker);
Ok(())
}
diff --git a/substrate/frame/staking-rewards/src/tests.rs b/substrate/frame/staking-rewards/src/tests.rs
index 1494f1c776b450c6b459e1c348713c1e54314945..ddcc8b5d29672be5120d5c96ded8d2d0a71d2123 100644
--- a/substrate/frame/staking-rewards/src/tests.rs
+++ b/substrate/frame/staking-rewards/src/tests.rs
@@ -55,6 +55,8 @@ fn pools() -> Vec<(u32, PoolInfo<u128, NativeOrWithId<u32>, u128, u64>)> {
}
mod create_pool {
+ use sp_runtime::traits::BadOrigin;
+
use super::*;
#[test]
@@ -168,7 +170,7 @@ mod create_pool {
}
#[test]
- fn non_existent_asset_fails() {
+ fn fails_for_non_existent_asset() {
new_test_ext().execute_with(|| {
let valid_asset = NativeOrWithId::<u32>::WithId(1);
let invalid_asset = NativeOrWithId::<u32>::WithId(200);
@@ -207,6 +209,27 @@ mod create_pool {
);
})
}
+
+ #[test]
+ fn fails_for_not_admin() {
+ new_test_ext().execute_with(|| {
+ let user = 100;
+ let staking_asset_id = NativeOrWithId::<u32>::Native;
+ let reward_asset_id = NativeOrWithId::<u32>::WithId(1);
+ let reward_rate_per_block = 100;
+ create_tokens(user, vec![reward_asset_id.clone()]);
+ assert_err!(
+ StakingRewards::create_pool(
+ RuntimeOrigin::signed(user),
+ Box::new(staking_asset_id.clone()),
+ Box::new(reward_asset_id.clone()),
+ reward_rate_per_block,
+ Some(999)
+ ),
+ BadOrigin
+ );
+ });
+ }
}
mod stake {
@@ -220,7 +243,6 @@ mod stake {
let staking_asset_id = NativeOrWithId::<u32>::WithId(1);
let reward_asset_id = NativeOrWithId::<u32>::Native;
let reward_rate_per_block = 100;
-
create_tokens(user, vec![staking_asset_id.clone()]);
assert_ok!(StakingRewards::create_pool(
@@ -258,17 +280,14 @@ mod stake {
}
#[test]
- fn non_existent_pool() {
+ fn fails_for_non_existent_pool() {
new_test_ext().execute_with(|| {
- // Setup
let user = 1;
let staking_asset_id = NativeOrWithId::<u32>::WithId(1);
-
create_tokens(user, vec![staking_asset_id.clone()]);
let non_existent_pool_id = 999;
- // User tries to stake tokens in a non-existent pool
assert_err!(
StakingRewards::stake(RuntimeOrigin::signed(user), non_existent_pool_id, 1000),
Error::<MockRuntime>::NonExistentPool
@@ -277,7 +296,101 @@ mod stake {
}
#[test]
- fn insufficient_balance() {
+ fn fails_for_insufficient_balance() {
// TODO: When we're able to freeze assets.
}
}
+
+mod unstake {
+ use super::*;
+
+ #[test]
+ fn success() {
+ new_test_ext().execute_with(|| {
+ // Setup
+ let user = 1;
+ let staking_asset_id = NativeOrWithId::<u32>::WithId(1);
+ let reward_asset_id = NativeOrWithId::<u32>::WithId(2);
+ let reward_rate_per_block = 100;
+ create_tokens(user, vec![staking_asset_id.clone(), reward_asset_id.clone()]);
+
+ assert_ok!(StakingRewards::create_pool(
+ RuntimeOrigin::signed(user),
+ Box::new(staking_asset_id.clone()),
+ Box::new(reward_asset_id.clone()),
+ reward_rate_per_block,
+ None
+ ));
+
+ let pool_id = 0;
+
+ // User stakes tokens
+ assert_ok!(StakingRewards::stake(RuntimeOrigin::signed(user), pool_id, 1000));
+
+ // User unstakes tokens
+ assert_ok!(StakingRewards::unstake(RuntimeOrigin::signed(user), pool_id, 500));
+
+ // Check that the user's staked amount is updated
+ assert_eq!(PoolStakers::<MockRuntime>::get(pool_id, user).unwrap().amount, 500);
+
+ // Check that the pool's total tokens staked is updated
+ assert_eq!(Pools::<MockRuntime>::get(pool_id).unwrap().total_tokens_staked, 500);
+
+ // User unstakes remaining tokens
+ assert_ok!(StakingRewards::unstake(RuntimeOrigin::signed(user), pool_id, 500));
+
+ // Check that the user's staked amount is zero
+ assert_eq!(PoolStakers::<MockRuntime>::get(pool_id, user).unwrap().amount, 0);
+
+ // Check that the pool's total tokens staked is zero
+ assert_eq!(Pools::<MockRuntime>::get(pool_id).unwrap().total_tokens_staked, 0);
+ });
+ }
+
+ #[test]
+ fn fails_for_non_existent_pool() {
+ new_test_ext().execute_with(|| {
+ // Setup
+ let user = 1;
+ let non_existent_pool_id = 999;
+
+ // User tries to unstake tokens from a non-existent pool
+ assert_err!(
+ StakingRewards::unstake(RuntimeOrigin::signed(user), non_existent_pool_id, 500),
+ Error::<MockRuntime>::NonExistentPool
+ );
+ });
+ }
+
+ #[test]
+ fn fails_for_insufficient_staked_amount() {
+ new_test_ext().execute_with(|| {
+ // Setup
+ let user = 1;
+ let staking_asset_id = NativeOrWithId::<u32>::WithId(1);
+ let reward_asset_id = NativeOrWithId::<u32>::WithId(2);
+ let reward_rate_per_block = 100;
+
+ create_tokens(user, vec![staking_asset_id.clone(), reward_asset_id.clone()]);
+
+ assert_ok!(StakingRewards::create_pool(
+ RuntimeOrigin::signed(user),
+ Box::new(staking_asset_id.clone()),
+ Box::new(reward_asset_id.clone()),
+ reward_rate_per_block,
+ None
+ ));
+
+ let pool_id = 0;
+
+ // User stakes tokens
+ assert_ok!(StakingRewards::stake(RuntimeOrigin::signed(user), pool_id, 1000));
+
+ // User tries to unstake more tokens than they have staked
+ assert_err!(
+ StakingRewards::unstake(RuntimeOrigin::signed(user), pool_id, 1500),
+ Error::<MockRuntime>::NotEnoughTokens
+ );
+ });
+ }
+}