diff --git a/substrate/frame/asset-rewards/src/lib.rs b/substrate/frame/asset-rewards/src/lib.rs index d9c1688a53b9de0f349b752311d9aa525d378c9f..35e1eaa58f3ace1f6aed9dcb02e6111e020bc5de 100644 --- a/substrate/frame/asset-rewards/src/lib.rs +++ b/substrate/frame/asset-rewards/src/lib.rs @@ -280,6 +280,8 @@ pub mod pallet { NotEnoughTokens, /// An operation was attempted on a non-existent pool. NonExistentPool, + /// An operation was attempted on a non-existent pool. + NonExistentStaker, /// An operation was attempted using a non-existent asset. NonExistentAsset, /// There was an error converting a block number. @@ -450,7 +452,8 @@ pub mod pallet { // Always start by updating the pool and staker rewards. let pool_info = Pools::<T>::get(pool_id).ok_or(Error::<T>::NonExistentPool)?; - let staker_info = PoolStakers::<T>::get(pool_id, &staker).unwrap_or_default(); + let staker_info = + PoolStakers::<T>::get(pool_id, &staker).ok_or(Error::<T>::NonExistentStaker)?; let (pool_info, mut staker_info) = Self::update_pool_and_staker_rewards(pool_info, staker_info)?; diff --git a/substrate/frame/asset-rewards/src/tests.rs b/substrate/frame/asset-rewards/src/tests.rs index 3f877cadc66a95b1cde45552ce4be70adba986f0..2bd405c6691874dfe4fd2d04fa4a4179519a9140 100644 --- a/substrate/frame/asset-rewards/src/tests.rs +++ b/substrate/frame/asset-rewards/src/tests.rs @@ -406,6 +406,139 @@ mod unstake { } } +mod harvest_rewards { + use super::*; + + #[test] + fn success() { + new_test_ext().execute_with(|| { + let staker = 1; + let pool_id = 0; + let reward_asset_id = NativeOrWithId::<u32>::Native; + create_default_pool(); + let pool_account_id = StakingRewards::pool_account_id(&pool_id).unwrap(); + <<MockRuntime as Config>::Assets>::set_balance( + reward_asset_id.clone(), + &pool_account_id, + 100_000, + ); + + // Stake + System::set_block_number(10); + assert_ok!(StakingRewards::stake(RuntimeOrigin::signed(staker), pool_id, 1000)); + + // Harvest + System::set_block_number(20); + let balance_before: <MockRuntime as Config>::Balance = + <<MockRuntime as Config>::Assets>::balance(reward_asset_id.clone(), &staker); + assert_ok!(StakingRewards::harvest_rewards( + RuntimeOrigin::signed(staker), + pool_id, + None + )); + let balance_after = + <<MockRuntime as Config>::Assets>::balance(reward_asset_id.clone(), &staker); + + // Assert + assert_eq!( + balance_after - balance_before, + 10 * Pools::<MockRuntime>::get(pool_id).unwrap().reward_rate_per_block + ); + assert_eq!( + *events().last().unwrap(), + Event::<MockRuntime>::RewardsHarvested { + who: staker, + staker, + pool_id, + amount: 10 * Pools::<MockRuntime>::get(pool_id).unwrap().reward_rate_per_block + } + ); + }); + } + + #[test] + fn succeeds_when_harvesting_for_other_staker() { + new_test_ext().execute_with(|| { + let staker = 1; + let harvester = 2; + let pool_id = 0; + let reward_asset_id = NativeOrWithId::<u32>::Native; + create_default_pool(); + let pool_account_id = StakingRewards::pool_account_id(&pool_id).unwrap(); + <<MockRuntime as Config>::Assets>::set_balance( + reward_asset_id.clone(), + &pool_account_id, + 100_000, + ); + + // Stake + System::set_block_number(10); + assert_ok!(StakingRewards::stake(RuntimeOrigin::signed(staker), pool_id, 1000)); + System::set_block_number(20); + + // Harvest + let balance_before: <MockRuntime as Config>::Balance = + <<MockRuntime as Config>::Assets>::balance(reward_asset_id.clone(), &staker); + assert_ok!(StakingRewards::harvest_rewards( + RuntimeOrigin::signed(harvester), + pool_id, + Some(staker) + )); + let balance_after = + <<MockRuntime as Config>::Assets>::balance(reward_asset_id.clone(), &staker); + + // Assert + assert_eq!( + balance_after - balance_before, + 10 * Pools::<MockRuntime>::get(pool_id).unwrap().reward_rate_per_block + ); + assert_eq!( + *events().last().unwrap(), + Event::<MockRuntime>::RewardsHarvested { + who: harvester, + staker, + pool_id, + amount: 10 * Pools::<MockRuntime>::get(pool_id).unwrap().reward_rate_per_block + } + ); + }); + } + + #[test] + fn fails_for_non_existent_staker() { + new_test_ext().execute_with(|| { + let non_existent_staker = 999; + + create_default_pool(); + assert_err!( + StakingRewards::harvest_rewards( + RuntimeOrigin::signed(non_existent_staker), + 0, + None + ), + Error::<MockRuntime>::NonExistentStaker + ); + }); + } + + #[test] + fn fails_for_non_existent_pool() { + new_test_ext().execute_with(|| { + let staker = 1; + let non_existent_pool_id = 999; + + assert_err!( + StakingRewards::harvest_rewards( + RuntimeOrigin::signed(staker), + non_existent_pool_id, + None + ), + Error::<MockRuntime>::NonExistentPool + ); + }); + } +} + mod set_pool_admin { use super::*;