Update reward pool on provider_boost or unstake #1699

pallets/capacity/src/benchmarking.rs

@@ -37,14 +37,67 @@ pub fn set_up_epoch<T: Config>(current_block: BlockNumberFor<T>, current_epoch:
 	CurrentEpochInfo::<T>::set(EpochInfo { epoch_start });
 }
 
+pub fn set_era_and_reward_pool_at_block<T: Config>(
+	era_index: T::RewardEra,
+	started_at: BlockNumberFor<T>,
+	total_staked_token: BalanceOf<T>,
+) {
+	let era_info: RewardEraInfo<T::RewardEra, BlockNumberFor<T>> =
+		RewardEraInfo { era_index, started_at };
+	let total_reward_pool: BalanceOf<T> =
+		T::MinimumStakingAmount::get().saturating_add(1_100u32.into());
+	CurrentEraInfo::<T>::set(era_info);
+	let pool_info: RewardPoolInfo<BalanceOf<T>> = RewardPoolInfo {
+		total_staked_token,
+		total_reward_pool,
+		unclaimed_balance: total_reward_pool,
+	};
+	StakingRewardPool::<T>::insert(era_index, pool_info);
+}
+
+// caller stakes the given amount to the given target
+pub fn setup_provider_stake<T: Config>(
+	caller: &T::AccountId,
+	target: &MessageSourceId,
+	staking_amount: BalanceOf<T>,
+) {
+	let capacity_amount: BalanceOf<T> = Capacity::<T>::capacity_generated(staking_amount);
+
+	let mut staking_account = StakingDetails::<T>::default();
+	let mut target_details = StakingTargetDetails::<BalanceOf<T>>::default();
+	let mut capacity_details =
+		CapacityDetails::<BalanceOf<T>, <T as Config>::EpochNumber>::default();
+
+	staking_account.deposit(staking_amount);
+	target_details.deposit(staking_amount, capacity_amount);
+	capacity_details.deposit(&staking_amount, &capacity_amount);
+
+	Capacity::<T>::set_staking_account_and_lock(caller, &staking_account)
+		.expect("Failed to set staking account");
+	Capacity::<T>::set_target_details_for(caller, *target, target_details);
+	Capacity::<T>::set_capacity_for(*target, capacity_details);
+}
+
+// fill up unlock chunks to max bound - 1
+fn fill_unlock_chunks<T: Config>(caller: &T::AccountId, count: u32) {
+	let mut unlocking: UnlockChunkList<T> = BoundedVec::default();
+	for _i in 0..count {
+		let unlock_chunk: UnlockChunk<BalanceOf<T>, T::EpochNumber> =
+			UnlockChunk { value: 1u32.into(), thaw_at: 3u32.into() };
+		assert_ok!(unlocking.try_push(unlock_chunk));
+	}
+	UnstakeUnlocks::<T>::set(caller, Some(unlocking));
+}
+
 benchmarks! {
 	stake {
 		let caller: T::AccountId = create_funded_account::<T>("account", SEED, 105u32);
 		let amount: BalanceOf<T> = T::MinimumStakingAmount::get();
 		let capacity: BalanceOf<T> = Capacity::<T>::capacity_generated(amount);
 		let target = 1;
-		let staking_type = StakingType::MaximumCapacity;
+		let staking_type = MaximumCapacity;
 
+		set_era_and_reward_pool_at_block::<T>(1u32.into(), 1u32.into(), 1_000u32.into());
 		register_provider::<T>(target, "Foo");
 
 	}: _ (RawOrigin::Signed(caller.clone()), target, amount)
@@ -57,12 +110,7 @@ benchmarks! {
 
 	withdraw_unstaked {
 		let caller: T::AccountId = create_funded_account::<T>("account", SEED, 5u32);
-		let mut unlocking: UnlockChunkList<T> = BoundedVec::default();
-		for _i in 0..T::MaxUnlockingChunks::get() {
-			let unlock_chunk: UnlockChunk<BalanceOf<T>, T::EpochNumber> = UnlockChunk { value: 1u32.into(), thaw_at: 3u32.into() };
-			assert_ok!(unlocking.try_push(unlock_chunk));
-		}
-		UnstakeUnlocks::<T>::set(&caller, Some(unlocking));
+		fill_unlock_chunks::<T>(&caller, T::MaxUnlockingChunks::get());
 
 		CurrentEpoch::<T>::set(T::EpochNumber::from(5u32));
 
@@ -90,31 +138,17 @@ benchmarks! {
 		let target = 1;
 		let block_number = 4u32;
 
-		let mut staking_account = StakingDetails::<T>::default();
-		let mut target_details = StakingTargetDetails::<BalanceOf<T>>::default();
-		let mut capacity_details = CapacityDetails::<BalanceOf<T>, <T as Config>::EpochNumber>::default();
-
-		staking_account.deposit(staking_amount);
-		target_details.deposit(staking_amount, capacity_amount);
-		capacity_details.deposit(&staking_amount, &capacity_amount);
-
-		Capacity::<T>::set_staking_account_and_lock(&caller.clone(), &staking_account).expect("Failed to set staking account");
-		Capacity::<T>::set_target_details_for(&caller.clone(), target, target_details);
-		Capacity::<T>::set_capacity_for(target, capacity_details);
-
-		// fill up unlock chunks to max bound - 1
-		let count = T::MaxUnlockingChunks::get()-1;
-		let mut unlocking: UnlockChunkList<T> = BoundedVec::default();
-		for _i in 0..count {
-			let unlock_chunk: UnlockChunk<BalanceOf<T>, T::EpochNumber> = UnlockChunk { value: 1u32.into(), thaw_at: 3u32.into() };
-			assert_ok!(unlocking.try_push(unlock_chunk));
-		}
-		UnstakeUnlocks::<T>::set(&caller, Some(unlocking));
-
-
+		set_era_and_reward_pool_at_block::<T>(1u32.into(), 1u32.into(), 1_000u32.into());
+		setup_provider_stake::<T>(&caller, &target, staking_amount);
+		fill_unlock_chunks::<T>(&caller, T::MaxUnlockingChunks::get() - 1);
 	}: _ (RawOrigin::Signed(caller.clone()), target, unstaking_amount.into())
 	verify {
-		assert_last_event::<T>(Event::<T>::UnStaked {account: caller, target: target, amount: unstaking_amount.into(), capacity: Capacity::<T>::calculate_capacity_reduction(unstaking_amount.into(), staking_amount, capacity_amount) }.into());
+		assert_last_event::<T>(Event::<T>::UnStaked {
+			account: caller.clone(),
+			target,
+			amount: unstaking_amount.into(),
+			capacity: Capacity::<T>::calculate_capacity_reduction(unstaking_amount.into(), staking_amount,capacity_amount)
+		}.into());
 	}
 
 	set_epoch_length {
@@ -125,7 +159,49 @@ benchmarks! {
 		assert_last_event::<T>(Event::<T>::EpochLengthUpdated {blocks: epoch_length}.into());
 	}
 
+	change_staking_target {
+		let caller: T::AccountId = create_funded_account::<T>("account", SEED, 5u32);
+		let from_msa = 33;
+		let to_msa = 34;
+		// amount in addition to minimum
+		let from_msa_amount: BalanceOf<T> = T::MinimumStakingAmount::get().saturating_add(31u32.into());
+		let to_msa_amount: BalanceOf<T> = T::MinimumStakingAmount::get().saturating_add(1u32.into());
+
+		set_era_and_reward_pool_at_block::<T>(1u32.into(), 1u32.into(), 1_000u32.into());
+		register_provider::<T>(from_msa, "frommsa");
+		register_provider::<T>(to_msa, "tomsa");
+		setup_provider_stake::<T>(&caller, &from_msa, from_msa_amount);
+		setup_provider_stake::<T>(&caller, &to_msa, to_msa_amount);
+		let restake_amount: BalanceOf<T> = from_msa_amount.saturating_sub(10u32.into());
+
+	}: _ (RawOrigin::Signed(caller.clone(), ), from_msa, to_msa, restake_amount)
+	verify {
+		assert_last_event::<T>(Event::<T>::StakingTargetChanged {
+			account: caller,
+			from_msa,
+			to_msa,
+			amount: restake_amount.into()
+		}.into());
+	}
+
+	provider_boost {
+		let caller: T::AccountId = create_funded_account::<T>("boostaccount", SEED, 260u32);
+		let boost_amount: BalanceOf<T> = T::MinimumStakingAmount::get().saturating_add(1u32.into());
+		let capacity: BalanceOf<T> = Capacity::<T>::capacity_generated(<T>::RewardsProvider::capacity_boost(boost_amount));
+		let target = 1;
+
+		set_era_and_reward_pool_at_block::<T>(1u32.into(), 1u32.into(), 1_000u32.into());
+		register_provider::<T>(target, "Foo");
+
+	}: _ (RawOrigin::Signed(caller.clone()), target, boost_amount)
+	verify {
+		assert!(StakingAccountLedger::<T>::contains_key(&caller));
+		assert!(StakingTargetLedger::<T>::contains_key(&caller, target));
+		assert!(CapacityLedger::<T>::contains_key(target));
+		assert_last_event::<T>(Event::<T>::ProviderBoosted {account: caller, amount: boost_amount, target, capacity}.into());
+	}
+
 	impl_benchmark_test_suite!(Capacity,
-		crate::tests::mock::new_test_ext(),
-		crate::tests::mock::Test);
+		tests::mock::new_test_ext(),
+		tests::mock::Test);
 }