add integration test

substrate/frame/asset-rewards/Cargo.toml

@@ -17,7 +17,7 @@ targets = ["x86_64-unknown-linux-gnu"]
 
 [dependencies]
 codec = { package = "parity-scale-codec", version = "3.6.1", default-features = false }
-frame-support = { path = "../support", default-features = false }
+frame-support = { path = "../support", default-features = false, features = ["experimental"] }
 frame-system = { path = "../system", default-features = false }
 frame-benchmarking = { path = "../benchmarking", default-features = false, optional = true }
 scale-info = { version = "2.10.0", default-features = false, features = ["derive"] }

substrate/frame/asset-rewards/src/lib.rs

@@ -80,7 +80,7 @@ pub type PoolId = u32;
 pub(crate) const PRECISION_SCALING_FACTOR: u32 = u32::MAX;
 
 /// A pool staker.
-#[derive(Default, Decode, Encode, MaxEncodedLen, TypeInfo)]
+#[derive(Debug, Default, Decode, Encode, MaxEncodedLen, TypeInfo)]
 pub struct PoolStakerInfo<Balance> {
 	/// Amount of tokens staked.
 	amount: Balance,
@@ -416,7 +416,7 @@ pub mod pallet {
 
 			// Emit event.
 			Self::deposit_event(Event::RewardsHarvested {
-				who: caller,
+				who: caller.clone(),
 				staker,
 				pool_id,
 				amount: staker_info.rewards,
@@ -424,6 +424,7 @@ pub mod pallet {
 
 			// Reset staker rewards.
 			staker_info.rewards = 0u32.into();
+			PoolStakers::<T>::insert(pool_id, &caller, staker_info);
 
 			Ok(())
 		}
@@ -479,7 +480,7 @@ pub mod pallet {
 
 	impl<T: Config> Pallet<T> {
 		/// Derive a pool account ID from the pallet's ID.
-		fn pool_account_id(id: &PoolId) -> Result<T::AccountId, DispatchError> {
+		pub fn pool_account_id(id: &PoolId) -> Result<T::AccountId, DispatchError> {
 			if Pools::<T>::contains_key(id) {
 				Ok(T::PalletId::get().into_sub_account_truncating(id))
 			} else {
@@ -488,18 +489,19 @@ pub mod pallet {
 		}
 
 		/// Update pool reward state.
-		fn update_pool_rewards(pool_id: &PoolId, staker: &T::AccountId) -> DispatchResult {
+		pub fn update_pool_rewards(pool_id: &PoolId, staker: &T::AccountId) -> DispatchResult {
 			let reward_per_token = Self::reward_per_token(pool_id)?;
 
-			let mut pool_info = Pools::<T>::get(pool_id).ok_or(Error::<T>::NonExistentPool)?;
-			pool_info.last_update_block = frame_system::Pallet::<T>::block_number();
-			Pools::<T>::insert(pool_id, pool_info);
-
 			let mut staker_info = PoolStakers::<T>::get(pool_id, staker).unwrap_or_default();
 			staker_info.rewards = Self::derive_rewards(pool_id, staker)?;
 			staker_info.reward_per_token_paid = reward_per_token;
 			PoolStakers::<T>::insert(pool_id, staker, staker_info);
 
+			let mut pool_info = Pools::<T>::get(pool_id).ok_or(Error::<T>::NonExistentPool)?;
+			pool_info.last_update_block = frame_system::Pallet::<T>::block_number();
+			pool_info.reward_per_token_stored = reward_per_token;
+			Pools::<T>::insert(pool_id, pool_info);
+
 			Ok(())
 		}
 
@@ -510,7 +512,7 @@ pub mod pallet {
 			let pool_info = Pools::<T>::get(pool_id).ok_or(Error::<T>::NonExistentPool)?;
 
 			if pool_info.total_tokens_staked.eq(&0u32.into()) {
-				return Ok(0u32.into());
+				return Ok(pool_info.reward_per_token_stored)
 			}
 
 			let blocks_elapsed: u32 = match frame_system::Pallet::<T>::block_number()
@@ -521,15 +523,13 @@ pub mod pallet {
 				Err(_) => return Err(Error::<T>::BlockNumberConversionError.into()),
 			};
 
-			Ok(pool_info
-				.reward_per_token_stored
-				.saturating_add(
-					pool_info
-						.reward_rate_per_block
-						.saturating_mul(blocks_elapsed.into())
-						.saturating_mul(PRECISION_SCALING_FACTOR.into()),
-				)
-				.ensure_div(pool_info.total_tokens_staked)?)
+			Ok(pool_info.reward_per_token_stored.saturating_add(
+				pool_info
+					.reward_rate_per_block
+					.saturating_mul(blocks_elapsed.into())
+					.saturating_mul(PRECISION_SCALING_FACTOR.into())
+					.ensure_div(pool_info.total_tokens_staked)?,
+			))
 		}
 
 		/// Derives the amount of rewards earned by a staker.

substrate/frame/asset-rewards/src/mock.rs

@@ -157,7 +157,7 @@ pub(crate) fn new_test_ext() -> sp_io::TestExternalities {
 	// .unwrap();
 
 	pallet_balances::GenesisConfig::<MockRuntime> {
-		balances: vec![(1, 10000), (2, 20000), (3, 30000), (4, 40000)],
+		balances: vec![(1, 10000), (2, 20000), (3, 30000), (4, 40000), (10, 40000), (20, 40000)],
 	}
 	.assimilate_storage(&mut t)
 	.unwrap();

substrate/frame/asset-rewards/src/tests.rs

@@ -16,13 +16,10 @@
 // limitations under the License.
 
 use crate::{mock::*, *};
-use frame_support::{assert_err, assert_ok, traits::fungible::NativeOrWithId};
+use frame_support::{assert_err, assert_ok, hypothetically, traits::fungible::NativeOrWithId};
 
 fn create_tokens(owner: u128, tokens: Vec<NativeOrWithId<u32>>) {
-	create_tokens_with_ed(owner, tokens, 1)
-}
-
-fn create_tokens_with_ed(owner: u128, tokens: Vec<NativeOrWithId<u32>>, ed: u128) {
+	let ed = 1;
 	for token_id in tokens {
 		let asset_id = match token_id {
 			NativeOrWithId::WithId(id) => id,
@@ -412,3 +409,138 @@ mod unstake {
 		});
 	}
 }
+
+mod integration {
+	use super::*;
+
+	/// Assert that an amount has been hypothetically earned by a staker.
+	fn assert_hypothetically_earned(
+		staker: u128,
+		expected_earned: u128,
+		pool_id: u32,
+		reward_asset_id: NativeOrWithId<u32>,
+	) {
+		hypothetically!({
+			// Get the pre-harvest balance.
+			let balance_before: <MockRuntime as Config>::Balance =
+				<<MockRuntime as Config>::Assets>::balance(reward_asset_id.clone(), &staker);
+
+			// Harvest the rewards.
+			assert_ok!(StakingRewards::harvest_rewards(
+				RuntimeOrigin::signed(staker),
+				pool_id,
+				None
+			));
+
+			// Sanity check: staker rewards are reset to 0.
+			assert_eq!(PoolStakers::<MockRuntime>::get(pool_id, staker).unwrap().rewards, 0);
+
+			// Check that the staker has earned the expected amount.
+			let balance_after =
+				<<MockRuntime as Config>::Assets>::balance(reward_asset_id.clone(), &staker);
+			assert_eq!(
+				balance_after - balance_before,
+				<u128 as Into<<MockRuntime as Config>::Balance>>::into(expected_earned)
+			);
+		});
+	}
+
+	#[test]
+	/// In this integration test scenario, we will consider 2 stakers each staking and unstaking at
+	/// different intervals, and assert their claimable rewards are as expected.
+	///
+	/// Note: There are occasionally off by 1 errors due to rounding. In practice, this is
+	/// insignificant.
+	fn two_stakers() {
+		new_test_ext().execute_with(|| {
+			// Setup
+			let admin = 1;
+			let staker1 = 10u128;
+			let staker2 = 20;
+			let staking_asset_id = NativeOrWithId::<u32>::WithId(1);
+			let reward_asset_id = NativeOrWithId::<u32>::Native;
+			let reward_rate_per_block = 100;
+			create_tokens(admin, vec![staking_asset_id.clone()]);
+			assert_ok!(StakingRewards::create_pool(
+				RuntimeOrigin::signed(admin),
+				Box::new(staking_asset_id.clone()),
+				Box::new(reward_asset_id.clone()),
+				reward_rate_per_block,
+				None
+			));
+			let pool_id = 0;
+			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,
+			);
+			<<MockRuntime as Config>::Assets>::set_balance(
+				staking_asset_id.clone(),
+				&staker1,
+				100_000,
+			);
+			<<MockRuntime as Config>::Assets>::set_balance(
+				staking_asset_id.clone(),
+				&staker2,
+				100_000,
+			);
+
+			// Block 7: Staker 1 stakes 100 tokens.
+			System::set_block_number(7);
+			assert_ok!(StakingRewards::stake(RuntimeOrigin::signed(staker1), pool_id, 100));
+			// At this point
+			// - Staker 1 has earned 0 tokens.
+			// - Staker 1 is earning 100 tokens per block.
+
+			// Check that Staker 1 has earned 0 tokens.
+			assert_hypothetically_earned(staker1, 0, pool_id, reward_asset_id.clone());
+
+			// Block 9: Staker 2 stakes 100 tokens.
+			System::set_block_number(9);
+			assert_ok!(StakingRewards::stake(RuntimeOrigin::signed(staker2), pool_id, 100));
+			// At this point
+			// - Staker 1 has earned 200 (100*2) tokens.
+			// - Staker 2 has earned 0 tokens.
+			// - Staker 1 is earning 50 tokens per block.
+			// - Staker 2 is earning 50 tokens per block.
+
+			// Check that Staker 1 has earned 200 tokens and Staker 2 has earned 0 tokens.
+			assert_hypothetically_earned(staker1, 200, pool_id, reward_asset_id.clone());
+			assert_hypothetically_earned(staker2, 0, pool_id, reward_asset_id.clone());
+
+			// Block 12: Staker 1 stakes an additional 100 tokens.
+			System::set_block_number(12);
+			assert_ok!(StakingRewards::stake(RuntimeOrigin::signed(staker1), pool_id, 100));
+			// At this point
+			// - Staker 1 has earned 350 (200 + (50 * 3)) tokens.
+			// - Staker 2 has earned 150 (50 * 3) tokens.
+			// - Staker 1 is earning 66.66 tokens per block.
+			// - Staker 2 is earning 33.33 tokens per block.
+
+			// Check that Staker 1 has earned 350 tokens and Staker 2 has earned 150 tokens.
+			assert_hypothetically_earned(staker1, 349, pool_id, reward_asset_id.clone());
+			assert_hypothetically_earned(staker2, 149, pool_id, reward_asset_id.clone());
+
+			// Block 22: Staker 1 unstakes 100 tokens.
+			System::set_block_number(22);
+			assert_ok!(StakingRewards::unstake(RuntimeOrigin::signed(staker1), pool_id, 100));
+			// - Staker 1 has earned 1016 (350 + 66.66 * 10) tokens.
+			// - Staker 2 has earned 483  (150 + 33.33 * 10) tokens.
+			// - Staker 1 is earning 50 tokens per block.
+			// - Staker 2 is earning 50 tokens per block.
+			assert_hypothetically_earned(staker1, 1015, pool_id, reward_asset_id.clone());
+			assert_hypothetically_earned(staker2, 483, pool_id, reward_asset_id.clone());
+
+			// Block 23: Staker 1 unstakes 100 tokens.
+			System::set_block_number(23);
+			assert_ok!(StakingRewards::unstake(RuntimeOrigin::signed(staker1), pool_id, 100));
+			// - Staker 1 has earned 1065 (1015 + 50) tokens.
+			// - Staker 2 has earned 533  (483 + 50) tokens.
+			// - Staker 1 is earning 0 tokens per block.
+			// - Staker 2 is earning 100 tokens per block.
+			assert_hypothetically_earned(staker1, 1064, pool_id, reward_asset_id.clone());
+			assert_hypothetically_earned(staker2, 533, pool_id, reward_asset_id.clone());
+		});
+	}
+}