From a5f34edcf8b37848f384810130d5b44ef9004e29 Mon Sep 17 00:00:00 2001
From: Shannon Wells <shannonwells@users.noreply.github.com>
Date: Thu, 25 Apr 2024 14:54:08 -0700
Subject: [PATCH] Update reward pool on `provider_boost` or `unstake` #1699
 (#1948)

The goal of this PR is to update the StakingRewardPool on a `provider_boost` or `unstake` extrinsic call.

Closes #1699

Also includes a bunch of lint changes

- [x] Design doc(s) updated
- [x] Tests added
- [x] Benchmarks corrected and updated
- [x] Weights updated
---
 pallets/capacity/src/benchmarking.rs          | 140 +++++++++++++----
 pallets/capacity/src/lib.rs                   |  70 +++++----
 pallets/capacity/src/migration/mod.rs         |   4 +-
 .../src/migration/provider_boost_init.rs      |  83 +++++------
 .../src/tests/change_staking_target_tests.rs  |  35 +----
 pallets/capacity/src/tests/eras_tests.rs      |  85 +++++------
 pallets/capacity/src/tests/mock.rs            |  17 ++-
 .../src/tests/provider_boost_tests.rs         |  47 ++++--
 .../capacity/src/tests/replenishment_tests.rs |   4 +-
 .../src/tests/rewards_provider_tests.rs       |   7 +-
 .../src/tests/stake_and_deposit_tests.rs      |   6 +-
 pallets/capacity/src/tests/testing_utils.rs   |  18 ++-
 pallets/capacity/src/tests/unstaking_tests.rs |  55 ++++++-
 .../src/tests/withdraw_unstaked_tests.rs      |   6 +-
 pallets/capacity/src/types.rs                 |  23 ++-
 pallets/capacity/src/weights.rs               | 141 +++++++++++++-----
 16 files changed, 494 insertions(+), 247 deletions(-)

diff --git a/pallets/capacity/src/benchmarking.rs b/pallets/capacity/src/benchmarking.rs
index a0ab1d5383..014e81fa11 100644
--- a/pallets/capacity/src/benchmarking.rs
+++ b/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);
 }
diff --git a/pallets/capacity/src/lib.rs b/pallets/capacity/src/lib.rs
index 179989f5fb..c025451f3a 100644
--- a/pallets/capacity/src/lib.rs
+++ b/pallets/capacity/src/lib.rs
@@ -90,8 +90,8 @@ pub mod weights;
 pub(crate) type BalanceOf<T> =
 	<<T as Config>::Currency as InspectFungible<<T as frame_system::Config>::AccountId>>::Balance;
 
-use frame_system::pallet_prelude::*;
 use crate::StakingType::{MaximumCapacity, ProviderBoost};
+use frame_system::pallet_prelude::*;
 
 #[frame_support::pallet]
 pub mod pallet {
@@ -269,13 +269,13 @@ pub mod pallet {
 	#[pallet::whitelist_storage]
 	#[pallet::getter(fn get_current_era)]
 	pub type CurrentEraInfo<T: Config> =
-	StorageValue<_, RewardEraInfo<T::RewardEra, BlockNumberFor<T>>, ValueQuery>;
+		StorageValue<_, RewardEraInfo<T::RewardEra, BlockNumberFor<T>>, ValueQuery>;
 
 	/// Reward Pool history
 	#[pallet::storage]
 	#[pallet::getter(fn get_reward_pool_for_era)]
 	pub type StakingRewardPool<T: Config> =
-	CountedStorageMap<_, Twox64Concat, T::RewardEra, RewardPoolInfo<BalanceOf<T>>>;
+		CountedStorageMap<_, Twox64Concat, T::RewardEra, RewardPoolInfo<BalanceOf<T>>>;
 
 	// TODO: storage for staking history
 
@@ -405,16 +405,14 @@ pub mod pallet {
 		/// Too many change_staking_target calls made in this RewardEra. (20)
 		MaxRetargetsExceeded,
 		/// There are no unstaked token amounts that have passed their thaw period.
-		NoThawedTokenAvailable
+		NoThawedTokenAvailable,
 	}
 
 	#[pallet::hooks]
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
 		fn on_initialize(current: BlockNumberFor<T>) -> Weight {
 			Self::start_new_epoch_if_needed(current)
-				.saturating_add(
-					Self::start_new_reward_era_if_needed(current)
-			)
+				.saturating_add(Self::start_new_reward_era_if_needed(current))
 		}
 	}
 
@@ -494,10 +492,12 @@ pub mod pallet {
 
 			ensure!(requested_amount > Zero::zero(), Error::<T>::UnstakedAmountIsZero);
 
-			let (actual_amount, staking_type) = Self::decrease_active_staking_balance(&unstaker, requested_amount)?;
+			let (actual_amount, staking_type) =
+				Self::decrease_active_staking_balance(&unstaker, requested_amount)?;
 			Self::add_unlock_chunk(&unstaker, actual_amount)?;
 
-			let capacity_reduction = Self::reduce_capacity(&unstaker, target, actual_amount, staking_type)?;
+			let capacity_reduction =
+				Self::reduce_capacity(&unstaker, target, actual_amount, staking_type)?;
 
 			Self::deposit_event(Event::UnStaked {
 				account: unstaker,
@@ -528,7 +528,7 @@ pub mod pallet {
 		}
 
 		/// Sets the target of the staking capacity to a new target.
-		/// This adds a chunk to `StakingAccountDetails.stake_change_unlocking chunks`, up to `T::MaxUnlockingChunks`.
+		/// This adds a chunk to `StakingDetails.stake_change_unlocking chunks`, up to `T::MaxUnlockingChunks`.
 		/// The staked amount and Capacity generated by `amount` originally targeted to the `from` MSA Id is reassigned to the `to` MSA Id.
 		/// Does not affect unstaking process or additional stake amounts.
 		/// Changing a staking target to a Provider when Origin has nothing staked them will retain the staking type.
@@ -541,7 +541,7 @@ pub mod pallet {
 		/// - [`Error::InvalidTarget`] if `to` does not belong to a registered Provider.
 		/// - [`Error::MaxRetargetsExceeded`] if origin has reached the maximimum number of retargets for the current RewardEra.
 		#[pallet::call_index(4)]
-		#[pallet::weight(T::WeightInfo::unstake())]
+		#[pallet::weight(T::WeightInfo::change_staking_target())]
 		pub fn change_staking_target(
 			origin: OriginFor<T>,
 			from: MessageSourceId,
@@ -603,7 +603,6 @@ pub mod pallet {
 
 			Ok(())
 		}
-
 	}
 }
 
@@ -627,7 +626,10 @@ impl<T: Config> Pallet<T> {
 
 		let staking_details = Self::get_staking_account_for(&staker).unwrap_or_default();
 		if !staking_details.active.is_zero() {
-			ensure!(staking_details.staking_type.eq(&staking_type), Error::<T>::CannotChangeStakingType);
+			ensure!(
+				staking_details.staking_type.eq(&staking_type),
+				Error::<T>::CannotChangeStakingType
+			);
 		}
 
 		let stakable_amount = Self::get_stakable_amount_for(&staker, amount);
@@ -653,7 +655,7 @@ impl<T: Config> Pallet<T> {
 		amount: &BalanceOf<T>,
 	) -> Result<(StakingDetails<T>, BalanceOf<T>), DispatchError> {
 		// FIXME: if one is boosting additional amounts, this will fail.
-		let (mut staking_details, stakable_amount) = 
+		let (mut staking_details, stakable_amount) =
 			Self::ensure_can_stake(staker, *target, *amount, ProviderBoost)?;
 		staking_details.staking_type = ProviderBoost;
 		// TODO: update when boost history is implemented fully
@@ -692,9 +694,7 @@ impl<T: Config> Pallet<T> {
 		target: &MessageSourceId,
 		amount: &BalanceOf<T>,
 	) -> Result<BalanceOf<T>, DispatchError> {
-		staking_details
-			.deposit(*amount)
-			.ok_or(ArithmeticError::Overflow)?;
+		staking_details.deposit(*amount).ok_or(ArithmeticError::Overflow)?;
 
 		// get the capacity generated by a Provider Boost
 		let capacity = Self::capacity_generated(T::RewardsProvider::capacity_boost(*amount));
@@ -706,12 +706,16 @@ impl<T: Config> Pallet<T> {
 		let mut capacity_details = Self::get_capacity_for(target).unwrap_or_default();
 		capacity_details.deposit(amount, &capacity).ok_or(ArithmeticError::Overflow)?;
 
+		let era = Self::get_current_era().era_index;
+		let mut reward_pool =
+			Self::get_reward_pool_for_era(era).ok_or(Error::<T>::EraOutOfRange)?;
+		reward_pool.total_staked_token = reward_pool.total_staked_token.saturating_add(*amount);
+
 		// TODO:  add boost history record for era when boost history is implemented
-		// let era = Self::get_current_era().era_index;
 		Self::set_staking_account_and_lock(staker, staking_details)?;
 		Self::set_target_details_for(staker, *target, target_details);
 		Self::set_capacity_for(*target, capacity_details);
-
+		Self::set_reward_pool(era, &reward_pool);
 		Ok(capacity)
 	}
 
@@ -767,6 +771,10 @@ impl<T: Config> Pallet<T> {
 		CapacityLedger::<T>::insert(target, capacity_details);
 	}
 
+	fn set_reward_pool(era: <T>::RewardEra, new_reward_pool: &RewardPoolInfo<BalanceOf<T>>) {
+		StakingRewardPool::<T>::set(era, Some(new_reward_pool.clone()));
+	}
+
 	/// Decrease a staking account's active token and reap if it goes below the minimum.
 	/// Returns: actual amount unstaked, plus the staking type + StakingDetails,
 	/// since StakingDetails may be reaped and staking type must be used to calculate the
@@ -782,6 +790,12 @@ impl<T: Config> Pallet<T> {
 		let actual_unstaked_amount = staking_account.withdraw(amount)?;
 		let staking_type = staking_account.staking_type;
 		Self::set_staking_account(unstaker, &staking_account);
+
+		let era = Self::get_current_era().era_index;
+		let mut reward_pool =
+			Self::get_reward_pool_for_era(era).ok_or(Error::<T>::EraOutOfRange)?;
+		reward_pool.total_staked_token = reward_pool.total_staked_token.saturating_sub(amount);
+		Self::set_reward_pool(era, &reward_pool.clone());
 		Ok((actual_unstaked_amount, staking_type))
 	}
 
@@ -843,6 +857,7 @@ impl<T: Config> Pallet<T> {
 		Ok(amount_withdrawn)
 	}
 
+	#[allow(unused)]
 	fn get_thaw_at_epoch() -> <T as Config>::EpochNumber {
 		let current_epoch: T::EpochNumber = Self::get_current_epoch();
 		let thaw_period = T::UnstakingThawPeriod::get();
@@ -854,7 +869,7 @@ impl<T: Config> Pallet<T> {
 		unstaker: &T::AccountId,
 		target: MessageSourceId,
 		amount: BalanceOf<T>,
-		staking_type: StakingType
+		staking_type: StakingType,
 	) -> Result<BalanceOf<T>, DispatchError> {
 		let mut staking_target_details = Self::get_target_for(&unstaker, &target)
 			.ok_or(Error::<T>::StakerTargetRelationshipNotFound)?;
@@ -930,7 +945,8 @@ impl<T: Config> Pallet<T> {
 	}
 
 	fn start_new_reward_era_if_needed(current_block: BlockNumberFor<T>) -> Weight {
-		let current_era_info: RewardEraInfo<T::RewardEra, BlockNumberFor<T>> = Self::get_current_era(); // 1r
+		let current_era_info: RewardEraInfo<T::RewardEra, BlockNumberFor<T>> =
+			Self::get_current_era(); // 1r
 
 		if current_block.saturating_sub(current_era_info.started_at) >= T::EraLength::get().into() {
 			let new_era_info = RewardEraInfo {
@@ -938,7 +954,7 @@ impl<T: Config> Pallet<T> {
 				started_at: current_block,
 			};
 
-			let current_reward_pool_info =
+			let current_reward_pool =
 				Self::get_reward_pool_for_era(current_era_info.era_index).unwrap_or_default(); // 1r
 
 			let past_eras_max = T::StakingRewardsPastErasMax::get();
@@ -952,9 +968,9 @@ impl<T: Config> Pallet<T> {
 			CurrentEraInfo::<T>::set(new_era_info); // 1w
 
 			let total_reward_pool =
-				T::RewardsProvider::reward_pool_size(current_reward_pool_info.total_staked_token);
+				T::RewardsProvider::reward_pool_size(current_reward_pool.total_staked_token);
 			let new_reward_pool = RewardPoolInfo {
-				total_staked_token: current_reward_pool_info.total_staked_token,
+				total_staked_token: current_reward_pool.total_staked_token,
 				total_reward_pool,
 				unclaimed_balance: total_reward_pool,
 			};
@@ -1043,7 +1059,11 @@ impl<T: Config> Nontransferable for Pallet<T> {
 	}
 
 	/// Increase all totals for the MSA's CapacityDetails.
-	fn deposit(msa_id: MessageSourceId, token_amount: Self::Balance, capacity_amount: Self::Balance) -> Result<(), DispatchError> {
+	fn deposit(
+		msa_id: MessageSourceId,
+		token_amount: Self::Balance,
+		capacity_amount: Self::Balance,
+	) -> Result<(), DispatchError> {
 		let mut capacity_details =
 			Self::get_capacity_for(msa_id).ok_or(Error::<T>::TargetCapacityNotFound)?;
 		capacity_details.deposit(&token_amount, &capacity_amount);
diff --git a/pallets/capacity/src/migration/mod.rs b/pallets/capacity/src/migration/mod.rs
index 0ea1e97d0a..911a883ec8 100644
--- a/pallets/capacity/src/migration/mod.rs
+++ b/pallets/capacity/src/migration/mod.rs
@@ -1,6 +1,6 @@
+/// initial values for ProviderBoost storage
+pub mod provider_boost_init;
 /// migrations to v2
 pub mod v2;
 /// migrations to v3
 pub mod v3;
-/// initial values for ProviderBoost storage
-pub mod provider_boost_init;
diff --git a/pallets/capacity/src/migration/provider_boost_init.rs b/pallets/capacity/src/migration/provider_boost_init.rs
index a304622eb9..e1783ae2fd 100644
--- a/pallets/capacity/src/migration/provider_boost_init.rs
+++ b/pallets/capacity/src/migration/provider_boost_init.rs
@@ -1,56 +1,53 @@
 use crate::{Config, CurrentEraInfo, RewardEraInfo, RewardPoolInfo, StakingRewardPool};
 use frame_support::{
-    pallet_prelude::Weight,
-    traits::OnRuntimeUpgrade
-};
-use frame_support::{
-    traits::Get,
+	pallet_prelude::Weight,
+	traits::{Get, OnRuntimeUpgrade},
 };
 
 #[cfg(feature = "try-runtime")]
 use sp_std::vec::Vec;
 
-/// Initialization during runtime upgrade for Provider Boost storage 
+/// Initialization during runtime upgrade for Provider Boost storage
 pub struct ProviderBoostInit<T: Config>(sp_std::marker::PhantomData<T>);
 
 impl<T: Config> OnRuntimeUpgrade for ProviderBoostInit<T> {
-    fn on_runtime_upgrade() -> Weight {
-        let current_era_info = CurrentEraInfo::<T>::get();
-        if current_era_info.eq(&RewardEraInfo::default()) { // 1r
-            let current_block = frame_system::Pallet::<T>::block_number(); // 1r
-            let era_index: T::RewardEra = 0u32.into();
-            CurrentEraInfo::<T>::set(RewardEraInfo { era_index, started_at: current_block, }); // 1w 
-            StakingRewardPool::<T>::insert(era_index, RewardPoolInfo::default());  // 1w
-            T::DbWeight::get().reads_writes(2, 2)
-        }   else {
-            T::DbWeight::get().reads(1)
-        } 
-    }
+	fn on_runtime_upgrade() -> Weight {
+		let current_era_info = CurrentEraInfo::<T>::get();
+		if current_era_info.eq(&RewardEraInfo::default()) {
+			// 1r
+			let current_block = frame_system::Pallet::<T>::block_number(); // 1r
+			let era_index: T::RewardEra = 0u32.into();
+			CurrentEraInfo::<T>::set(RewardEraInfo { era_index, started_at: current_block }); // 1w
+			StakingRewardPool::<T>::insert(era_index, RewardPoolInfo::default()); // 1w
+			T::DbWeight::get().reads_writes(2, 2)
+		} else {
+			T::DbWeight::get().reads(1)
+		}
+	}
 
-    #[cfg(feature = "try-runtime")]
-    fn pre_upgrade() -> Result<Vec<u8>, TryRuntimeError> {
-        if CurrentEraInfo::<T>::exists() {
-            log::info!("CurrentEraInfo exists; initialization should be skipped.");
-        } else {
-            log::info!("CurrentEraInfo not found. Initialization should proceed.");
-        }
-        if StakingRewardPool::<T>::iter().count() == 0usize {
-            log::info!("StakingRewardPool will be updated with Era 0");
-        } else {
-            log::info!("StakingRewardPool has already been initialized.")
-        }
-        Ok(Vec::default())
-    }
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade() -> Result<Vec<u8>, TryRuntimeError> {
+		if CurrentEraInfo::<T>::exists() {
+			log::info!("CurrentEraInfo exists; initialization should be skipped.");
+		} else {
+			log::info!("CurrentEraInfo not found. Initialization should proceed.");
+		}
+		if StakingRewardPool::<T>::iter().count() == 0usize {
+			log::info!("StakingRewardPool will be updated with Era 0");
+		} else {
+			log::info!("StakingRewardPool has already been initialized.")
+		}
+		Ok(Vec::default())
+	}
 
-    #[cfg(feature = "try-runtime")]
-    fn post_upgrade(_state: Vec<u8>) -> Result<(), TryRuntimeError> {
-        assert!(CurrentEraInfo::<T>::exists());
-        let current_block = frame_system::Pallet::<T>::block_number();
-        let info = CurrentEraInfo::<T>::get() ;
-        assert_eq!(info.started_at, current_block);
-        log::info!("CurrentEraInfo.started_at is set to {:?}.", info.started_at);
-        assert_eq!(StakingRewardPool::<T>::iter().count(), 1);
-        Ok(())
-    }    
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade(_state: Vec<u8>) -> Result<(), TryRuntimeError> {
+		assert!(CurrentEraInfo::<T>::exists());
+		let current_block = frame_system::Pallet::<T>::block_number();
+		let info = CurrentEraInfo::<T>::get();
+		assert_eq!(info.started_at, current_block);
+		log::info!("CurrentEraInfo.started_at is set to {:?}.", info.started_at);
+		assert_eq!(StakingRewardPool::<T>::iter().count(), 1);
+		Ok(())
+	}
 }
-
diff --git a/pallets/capacity/src/tests/change_staking_target_tests.rs b/pallets/capacity/src/tests/change_staking_target_tests.rs
index b10133152a..803ea8d2b9 100644
--- a/pallets/capacity/src/tests/change_staking_target_tests.rs
+++ b/pallets/capacity/src/tests/change_staking_target_tests.rs
@@ -3,9 +3,7 @@ use super::{
 	testing_utils::{setup_provider, staking_events},
 };
 use crate::*;
-use common_primitives::{
-	msa::MessageSourceId,
-};
+use common_primitives::msa::MessageSourceId;
 use frame_support::{assert_noop, assert_ok, traits::Get};
 
 // staker is unused unless amount > 0
@@ -28,14 +26,8 @@ fn assert_capacity_details(
 	assert_eq!(from_capacity_details, expected_from_details);
 }
 
-fn assert_target_details(
-	staker: u64,
-	msa_id: MessageSourceId,
-	amount: u64,
-	capacity: u64,
-) {
-	let expected_from_target_details: TestTargetDetails =
-		StakingTargetDetails { amount, capacity };
+fn assert_target_details(staker: u64, msa_id: MessageSourceId, amount: u64, capacity: u64) {
+	let expected_from_target_details: TestTargetDetails = StakingTargetDetails { amount, capacity };
 	let from_target_details = Capacity::get_target_for(staker, msa_id).unwrap();
 	assert_eq!(from_target_details, expected_from_target_details);
 }
@@ -80,19 +72,9 @@ fn do_retarget_flip_flop() {
 
 		for i in 0..4 {
 			if i % 2 == 0 {
-				assert_ok!(Capacity::do_retarget(
-					&staker,
-					&from_msa,
-					&to_msa,
-					&to_amount,
-				));
+				assert_ok!(Capacity::do_retarget(&staker, &from_msa, &to_msa, &to_amount,));
 			} else {
-				assert_ok!(Capacity::do_retarget(
-					&staker,
-					&to_msa,
-					&from_msa,
-					&to_amount,
-				));
+				assert_ok!(Capacity::do_retarget(&staker, &to_msa, &from_msa, &to_amount,));
 			}
 		}
 		assert_capacity_details(from_msa, 3, 600, 3);
@@ -210,10 +192,8 @@ fn do_retarget_deletes_staking_target_details_if_zero_balance() {
 
 		assert!(Capacity::get_target_for(staker, from_msa).is_none());
 
-		let expected_to_target_details: TestTargetDetails = StakingTargetDetails {
-			amount: 2 * amount,
-			capacity: 2,
-		};
+		let expected_to_target_details: TestTargetDetails =
+			StakingTargetDetails { amount: 2 * amount, capacity: 2 };
 		let to_target_details = Capacity::get_target_for(staker, to_msa).unwrap();
 		assert_eq!(to_target_details, expected_to_target_details);
 
@@ -288,6 +268,7 @@ fn change_staking_target_garbage_collects_thawed_chunks() {
 		let staking_account = 200u64;
 		let from_target: MessageSourceId = 3;
 		let to_target: MessageSourceId = 4;
+
 		setup_provider(&staking_account, &from_target, &staked_amount, ProviderBoost);
 		setup_provider(&staking_account, &to_target, &staked_amount, ProviderBoost);
 
diff --git a/pallets/capacity/src/tests/eras_tests.rs b/pallets/capacity/src/tests/eras_tests.rs
index b6ec94b0e9..f9fdcadbcb 100644
--- a/pallets/capacity/src/tests/eras_tests.rs
+++ b/pallets/capacity/src/tests/eras_tests.rs
@@ -1,22 +1,16 @@
-use super::{
-	mock::*,
-	testing_utils::{run_to_block, system_run_to_block},
+use super::mock::*;
+use crate::{
+	tests::testing_utils::{run_to_block, setup_provider, system_run_to_block},
+	Config, CurrentEraInfo, RewardEraInfo, RewardPoolInfo, StakingRewardPool,
+	StakingType::*,
 };
-use crate::{Config, CurrentEraInfo, RewardEraInfo, RewardPoolInfo, StakingRewardPool};
+use common_primitives::msa::MessageSourceId;
+use frame_support::assert_ok;
 use sp_core::Get;
 
 #[test]
-fn start_new_era_if_needed_updates_era_info_and_limits_reward_pool_size() {
+fn start_new_era_if_needed_updates_era_info() {
 	new_test_ext().execute_with(|| {
-		CurrentEraInfo::<Test>::set(RewardEraInfo { era_index: 1, started_at: 0 });
-		StakingRewardPool::<Test>::insert(
-			1,
-			RewardPoolInfo {
-				total_staked_token: 10_000,
-				total_reward_pool: 1_000,
-				unclaimed_balance: 1_000,
-			},
-		);
 		system_run_to_block(9);
 		for i in 1..4 {
 			let block_decade = i * 10;
@@ -25,8 +19,11 @@ fn start_new_era_if_needed_updates_era_info_and_limits_reward_pool_size() {
 			let current_era_info = CurrentEraInfo::<Test>::get();
 
 			let expected_era = i + 1;
-			assert_eq!(current_era_info.era_index, expected_era);
-			assert_eq!(current_era_info.started_at, block_decade);
+			assert_eq!(
+				current_era_info,
+				RewardEraInfo { era_index: expected_era, started_at: block_decade }
+			);
+
 			let past_eras_max: u32 = <Test as Config>::StakingRewardsPastErasMax::get();
 			assert!(StakingRewardPool::<Test>::count().le(&past_eras_max));
 			system_run_to_block(block_decade + 9);
@@ -34,41 +31,37 @@ fn start_new_era_if_needed_updates_era_info_and_limits_reward_pool_size() {
 	})
 }
 
+// Test that additional stake is carried over to the next era's RewardPoolInfo.
 #[test]
 fn start_new_era_if_needed_updates_reward_pool() {
 	new_test_ext().execute_with(|| {
-		CurrentEraInfo::<Test>::set(RewardEraInfo { era_index: 1, started_at: 0 });
-		StakingRewardPool::<Test>::insert(
-			1,
-			RewardPoolInfo {
-				total_staked_token: 10_000,
-				total_reward_pool: 1_000,
-				unclaimed_balance: 1_000,
-			},
-		);
 		system_run_to_block(8);
+		let staker = 10_000;
+		let provider_msa: MessageSourceId = 1;
+		let stake_amount = 600u64;
 
-		// TODO: Provider boost, after staking updates reward pool info #1699
-		// let staker = 10_000;
-		// let provider_msa: MessageSourceId = 1;
-		// let stake_amount = 600u64;
-		// setup_provider(&staker, &provider_msa, &stake_amount, ProviderBoost);
-
-		system_run_to_block(9);
-		run_to_block(10);
-		assert_eq!(StakingRewardPool::<Test>::count(), 2);
-		let current_reward_pool_info = StakingRewardPool::<Test>::get(2).unwrap();
-		assert_eq!(
-			current_reward_pool_info,
-			RewardPoolInfo {
-				total_staked_token: 10_000,
-				total_reward_pool: 1_000,
-				unclaimed_balance: 1_000,
-			}
-		);
+		// set up initial stake.
+		setup_provider(&staker, &provider_msa, &stake_amount, ProviderBoost);
 
-		// TODO: after staking updates reward pool info #1699
-		// system_run_to_block(19);
-		// run_to_block(20);
+		for i in 1u32..4 {
+			let era = i + 1;
+			let final_block = i * 10;
+			system_run_to_block(final_block - 1);
+			run_to_block(final_block);
+			assert_eq!(StakingRewardPool::<Test>::count(), era);
+			assert_eq!(
+				StakingRewardPool::<Test>::get(era).unwrap(),
+				RewardPoolInfo {
+					total_staked_token: (stake_amount * i as u64),
+					total_reward_pool: 10_000,
+					unclaimed_balance: 10_000,
+				}
+			);
+			assert_ok!(Capacity::provider_boost(
+				RuntimeOrigin::signed(staker),
+				provider_msa,
+				stake_amount
+			));
+		}
 	});
 }
diff --git a/pallets/capacity/src/tests/mock.rs b/pallets/capacity/src/tests/mock.rs
index cc6bbf307a..deda9345ef 100644
--- a/pallets/capacity/src/tests/mock.rs
+++ b/pallets/capacity/src/tests/mock.rs
@@ -1,6 +1,9 @@
 use crate as pallet_capacity;
 
-use crate::{BalanceOf, StakingRewardClaim, StakingRewardsProvider};
+use crate::{
+	tests::testing_utils::set_era_and_reward_pool_at_block, BalanceOf, StakingRewardClaim,
+	StakingRewardsProvider,
+};
 use common_primitives::{
 	node::{AccountId, Hash, ProposalProvider},
 	schema::{SchemaId, SchemaValidator},
@@ -15,7 +18,7 @@ use sp_runtime::{
 	traits::{BlakeTwo256, Convert, IdentityLookup},
 	AccountId32, BuildStorage, DispatchError, Perbill,
 };
-use sp_std::ops::{Div, Mul};
+use sp_std::ops::Mul;
 
 type Block = frame_system::mocking::MockBlockU32<Test>;
 
@@ -140,8 +143,9 @@ impl StakingRewardsProvider<Test> for TestStakingRewardsProvider {
 	type RewardEra = TestRewardEra;
 	type Hash = Hash; // use what's in common_primitives::node
 
-	fn reward_pool_size(total_staked: BalanceOf<Test>) -> BalanceOf<Test> {
-		total_staked.div(10u64)
+	// To reflect new economic model behavior of having a constant RewardPool amount.
+	fn reward_pool_size(_total_staked: BalanceOf<Test>) -> BalanceOf<Test> {
+		10_000u64.into()
 	}
 
 	fn staking_reward_total(
@@ -216,6 +220,9 @@ pub fn new_test_ext() -> sp_io::TestExternalities {
 	.unwrap();
 
 	let mut ext = sp_io::TestExternalities::new(t);
-	ext.execute_with(|| System::set_block_number(1));
+	ext.execute_with(|| {
+		System::set_block_number(1);
+		set_era_and_reward_pool_at_block(1, 0, 0);
+	});
 	ext
 }
diff --git a/pallets/capacity/src/tests/provider_boost_tests.rs b/pallets/capacity/src/tests/provider_boost_tests.rs
index 4d9aadeafe..8e274c79d0 100644
--- a/pallets/capacity/src/tests/provider_boost_tests.rs
+++ b/pallets/capacity/src/tests/provider_boost_tests.rs
@@ -1,10 +1,7 @@
 use super::{mock::*, testing_utils::*};
-use crate::{StakingDetails, Error, Event, StakingType, FreezeReason};
-use common_primitives::{msa::MessageSourceId};
-use frame_support::{
-	assert_noop, assert_ok, traits::fungible::InspectFreeze
-};
-use crate::Config;
+use crate::{Config, Error, Event, FreezeReason, RewardPoolInfo, StakingDetails, StakingType};
+use common_primitives::msa::MessageSourceId;
+use frame_support::{assert_noop, assert_ok, traits::fungible::InspectFreeze};
 
 #[test]
 fn provider_boost_works() {
@@ -27,15 +24,19 @@ fn provider_boost_works() {
 		// Check that the capacity generated is correct. (5% of amount staked, since 10% is what's in the mock)
 		let capacity_details = Capacity::get_capacity_for(target).unwrap();
 		assert_eq!(capacity_details.total_capacity_issued, 1u64);
-		
+
 		let events = staking_events();
 		assert_eq!(
 			events.first().unwrap(),
 			&Event::ProviderBoosted { account, target, amount, capacity }
 		);
 
-		assert_eq!(<Test as Config>::Currency::balance_frozen(&FreezeReason::CapacityStaking.into(), &account),
-				   200u64
+		assert_eq!(
+			<Test as Config>::Currency::balance_frozen(
+				&FreezeReason::CapacityStaking.into(),
+				&account
+			),
+			200u64
 		);
 
 		let target_details = Capacity::get_target_for(account, target).unwrap();
@@ -44,7 +45,33 @@ fn provider_boost_works() {
 }
 
 #[test]
-fn provider_boost_updates_boost_account_details() {
+fn provider_boost_updates_reward_pool_history() {
+	// two accounts staking to the same target
+	new_test_ext().execute_with(|| {
+		let account1 = 600;
+		let account2 = 500;
+		let target: MessageSourceId = 1;
+		let amount1 = 500;
+		let amount2 = 200;
+		register_provider(target, String::from("Foo"));
+
+		assert_ok!(Capacity::provider_boost(RuntimeOrigin::signed(account1), target, amount1));
+		assert_ok!(Capacity::provider_boost(RuntimeOrigin::signed(account2), target, amount2));
+
+		let reward_pool = Capacity::get_reward_pool_for_era(1).unwrap();
+		assert_eq!(
+			reward_pool,
+			RewardPoolInfo {
+				total_staked_token: 700,
+				total_reward_pool: 10_000, // TODO: get the constant value and use it
+				unclaimed_balance: 10_000,
+			}
+		);
+	});
+}
+
+#[test]
+fn provider_boost_updates_staking_details() {
 	new_test_ext().execute_with(|| {
 		let account = 600;
 		let target: MessageSourceId = 1;
diff --git a/pallets/capacity/src/tests/replenishment_tests.rs b/pallets/capacity/src/tests/replenishment_tests.rs
index 02d04a0238..1a7d2cacee 100644
--- a/pallets/capacity/src/tests/replenishment_tests.rs
+++ b/pallets/capacity/src/tests/replenishment_tests.rs
@@ -63,7 +63,7 @@ fn impl_can_replenish_is_false_when_last_replenished_at_is_greater_or_equal_curr
 
 		assert_eq!(Capacity::can_replenish(target_msa_id), false);
 
-		CurrentEpoch::<Test>::set(1u32.into());
+		CurrentEpoch::<Test>::set(1);
 
 		assert_eq!(Capacity::can_replenish(target_msa_id), false);
 	});
@@ -84,7 +84,7 @@ fn impl_can_replenish_is_true_when_last_replenished_at_is_less_than_current_epoc
 			last_replenished_at,
 		);
 
-		CurrentEpoch::<Test>::set(3u32.into());
+		CurrentEpoch::<Test>::set(3);
 
 		assert_eq!(Capacity::can_replenish(target_msa_id), true);
 	});
diff --git a/pallets/capacity/src/tests/rewards_provider_tests.rs b/pallets/capacity/src/tests/rewards_provider_tests.rs
index 2ea1a31cc5..781b882ca0 100644
--- a/pallets/capacity/src/tests/rewards_provider_tests.rs
+++ b/pallets/capacity/src/tests/rewards_provider_tests.rs
@@ -1,8 +1,7 @@
 use super::mock::*;
 use crate::{
-	CurrentEraInfo, Error, RewardEraInfo, RewardPoolInfo, StakingDetails,
-	StakingRewardClaim, StakingRewardPool, StakingRewardsProvider,
-	StakingType::*
+	CurrentEraInfo, Error, RewardEraInfo, RewardPoolInfo, StakingDetails, StakingRewardClaim,
+	StakingRewardPool, StakingRewardsProvider, StakingType::*,
 };
 use frame_support::assert_err;
 
@@ -18,7 +17,7 @@ fn test_staking_reward_total_happy_path() {
 		let proof = H256::random();
 		let payload: StakingRewardClaim<Test> = StakingRewardClaim {
 			claimed_reward: 1,
-			staking_account_end_state: StakingDetails { active: 1, staking_type: MaximumCapacity},
+			staking_account_end_state: StakingDetails { active: 1, staking_type: MaximumCapacity },
 			from_era: 1,
 			to_era: 5,
 		};
diff --git a/pallets/capacity/src/tests/stake_and_deposit_tests.rs b/pallets/capacity/src/tests/stake_and_deposit_tests.rs
index b2da63755e..47e67e8018 100644
--- a/pallets/capacity/src/tests/stake_and_deposit_tests.rs
+++ b/pallets/capacity/src/tests/stake_and_deposit_tests.rs
@@ -1,9 +1,11 @@
 use super::{mock::*, testing_utils::*};
-use crate::{BalanceOf, CapacityDetails, Config, Error, Event, FreezeReason, StakingDetails};
+use crate::{
+	BalanceOf, CapacityDetails, Config, Error, Event, FreezeReason, StakingDetails,
+	StakingType::MaximumCapacity,
+};
 use common_primitives::{capacity::Nontransferable, msa::MessageSourceId};
 use frame_support::{assert_noop, assert_ok, traits::fungible::InspectFreeze};
 use sp_runtime::ArithmeticError;
-use crate::StakingType::MaximumCapacity;
 
 #[test]
 fn stake_works() {
diff --git a/pallets/capacity/src/tests/testing_utils.rs b/pallets/capacity/src/tests/testing_utils.rs
index 2ecdd901c7..81b3a9a6b1 100644
--- a/pallets/capacity/src/tests/testing_utils.rs
+++ b/pallets/capacity/src/tests/testing_utils.rs
@@ -4,7 +4,10 @@ use frame_support::{assert_ok, traits::Hooks};
 #[allow(unused)]
 use sp_runtime::traits::SignedExtension;
 
-use crate::{BalanceOf, CapacityDetails, Config, Event, StakingType};
+use crate::{
+	BalanceOf, CapacityDetails, Config, CurrentEraInfo, Event, RewardEraInfo, RewardPoolInfo,
+	StakingRewardPool, StakingType,
+};
 use common_primitives::msa::MessageSourceId;
 
 pub fn staking_events() -> Vec<Event<Test>> {
@@ -89,3 +92,16 @@ pub fn setup_provider(
 		assert_eq!(account_staking_type, staking_type);
 	}
 }
+
+// Currently the reward pool is a constant, however it could change in the future.
+pub fn set_era_and_reward_pool_at_block(era_index: u32, started_at: u32, total_staked_token: u64) {
+	let era_info = RewardEraInfo { era_index, started_at };
+	let total_reward_pool = 10_000u64;
+	CurrentEraInfo::<Test>::set(era_info);
+	let pool_info = RewardPoolInfo {
+		total_staked_token,
+		total_reward_pool,
+		unclaimed_balance: total_reward_pool,
+	};
+	StakingRewardPool::<Test>::insert(era_index, pool_info);
+}
diff --git a/pallets/capacity/src/tests/unstaking_tests.rs b/pallets/capacity/src/tests/unstaking_tests.rs
index c689153b7c..ec2d4a8dd4 100644
--- a/pallets/capacity/src/tests/unstaking_tests.rs
+++ b/pallets/capacity/src/tests/unstaking_tests.rs
@@ -1,7 +1,8 @@
 use super::{mock::*, testing_utils::*};
 use crate as pallet_capacity;
 use crate::{
-	CapacityDetails, FreezeReason, StakingDetails, StakingTargetDetails, StakingType, UnlockChunk,
+	CapacityDetails, FreezeReason, RewardPoolInfo, StakingDetails, StakingTargetDetails,
+	StakingType, UnlockChunk,
 };
 use common_primitives::msa::MessageSourceId;
 use frame_support::{
@@ -231,7 +232,6 @@ fn unstaking_everything_reaps_staking_account() {
 		register_provider(target, String::from("WithdrawUnst"));
 		assert_ok!(Capacity::stake(RuntimeOrigin::signed(staker), target, amount));
 
-
 		run_to_block(1);
 		// unstake everything
 		assert_ok!(Capacity::unstake(RuntimeOrigin::signed(staker), target, 20));
@@ -257,3 +257,54 @@ fn unstake_when_not_staking_to_target_errors() {
 		);
 	})
 }
+
+#[test]
+fn unstake_provider_boosted_target_adjusts_reward_pool_total() {
+	new_test_ext().execute_with(|| {
+		// two accounts staking to the same target
+		let account1 = 600;
+		let target: MessageSourceId = 1;
+		let amount1 = 500;
+		let unstake_amount = 200;
+		register_provider(target, String::from("Foo"));
+		run_to_block(5); // ensures Capacity::on_initialize is run
+
+		assert_ok!(Capacity::provider_boost(RuntimeOrigin::signed(account1), target, amount1));
+		assert_ok!(Capacity::unstake(RuntimeOrigin::signed(account1), target, unstake_amount));
+
+		let reward_pool = Capacity::get_reward_pool_for_era(1).unwrap();
+		assert_eq!(
+			reward_pool,
+			RewardPoolInfo {
+				total_staked_token: 300,
+				total_reward_pool: 10_000, // TODO: get the constant value and use it
+				unclaimed_balance: 10_000,
+			}
+		);
+	});
+}
+
+#[test]
+fn unstake_fills_up_common_unlock_for_any_target() {
+	new_test_ext().execute_with(|| {
+		let staker = 10_000;
+
+		let target1 = 1;
+		let target2 = 2;
+		register_provider(target1, String::from("Test Target"));
+		register_provider(target2, String::from("Test Target"));
+
+		assert_ok!(Capacity::provider_boost(RuntimeOrigin::signed(staker), target1, 1_000));
+		assert_ok!(Capacity::provider_boost(RuntimeOrigin::signed(staker), target2, 2_000));
+
+		// max unlock chunks in mock is 4
+		for _i in 0..2 {
+			assert_ok!(Capacity::unstake(RuntimeOrigin::signed(staker), target1, 50));
+			assert_ok!(Capacity::unstake(RuntimeOrigin::signed(staker), target2, 50));
+		}
+		assert_noop!(
+			Capacity::unstake(RuntimeOrigin::signed(staker), target1, 50),
+			Error::<Test>::MaxUnlockingChunksExceeded
+		);
+	})
+}
diff --git a/pallets/capacity/src/tests/withdraw_unstaked_tests.rs b/pallets/capacity/src/tests/withdraw_unstaked_tests.rs
index 2e89709931..c688a18dcb 100644
--- a/pallets/capacity/src/tests/withdraw_unstaked_tests.rs
+++ b/pallets/capacity/src/tests/withdraw_unstaked_tests.rs
@@ -90,7 +90,11 @@ fn withdraw_unstaked_cleans_up_storage_and_removes_all_locks_if_no_stake_left()
 		run_to_block(21);
 		assert_ok!(Capacity::withdraw_unstaked(RuntimeOrigin::signed(staker)));
 
-		assert_eq!(<Test as Config>::Currency::balance_frozen(&FreezeReason::CapacityStaking.into(),&staker),
+		assert_eq!(
+			<Test as Config>::Currency::balance_frozen(
+				&FreezeReason::CapacityStaking.into(),
+				&staker
+			),
 			0u64
 		);
 		assert!(Capacity::get_unstake_unlocking_for(&staker).is_none());
diff --git a/pallets/capacity/src/types.rs b/pallets/capacity/src/types.rs
index 6760fbfdd3..fffed318bd 100644
--- a/pallets/capacity/src/types.rs
+++ b/pallets/capacity/src/types.rs
@@ -4,10 +4,9 @@ use frame_support::{BoundedVec, EqNoBound, PartialEqNoBound, RuntimeDebugNoBound
 use parity_scale_codec::{Decode, Encode, EncodeLike, MaxEncodedLen};
 use scale_info::TypeInfo;
 use sp_runtime::{
-	traits::{CheckedAdd, CheckedSub, Saturating},
+	traits::{AtLeast32BitUnsigned, CheckedAdd, CheckedSub, Saturating},
 	RuntimeDebug,
 };
-use sp_runtime::traits::AtLeast32BitUnsigned;
 #[cfg(any(feature = "runtime-benchmarks", test))]
 use sp_std::vec::Vec;
 
@@ -245,10 +244,22 @@ pub fn unlock_chunks_from_vec<T: Config>(chunks: &Vec<(u32, u32)>) -> UnlockChun
 }
 
 /// The information needed to track a Reward Era
-#[derive( PartialEq, Eq, Clone, Copy, Default, PartialOrd, Encode, Decode, RuntimeDebug, TypeInfo, MaxEncodedLen)]
+#[derive(
+	PartialEq,
+	Eq,
+	Clone,
+	Copy,
+	Default,
+	PartialOrd,
+	Encode,
+	Decode,
+	RuntimeDebug,
+	TypeInfo,
+	MaxEncodedLen,
+)]
 pub struct RewardEraInfo<RewardEra, BlockNumber>
-	where
-		RewardEra: AtLeast32BitUnsigned + EncodeLike,
+where
+	RewardEra: AtLeast32BitUnsigned + EncodeLike,
 {
 	/// the index of this era
 	pub era_index: RewardEra,
@@ -262,7 +273,7 @@ pub struct RewardEraInfo<RewardEra, BlockNumber>
 /// The unclaimed_balance is initialized to total_reward_pool and deducted whenever a
 /// valid claim occurs.
 #[derive(
-PartialEq, Eq, Clone, Default, PartialOrd, Encode, Decode, RuntimeDebug, TypeInfo, MaxEncodedLen,
+	PartialEq, Eq, Clone, Default, PartialOrd, Encode, Decode, RuntimeDebug, TypeInfo, MaxEncodedLen,
 )]
 pub struct RewardPoolInfo<Balance> {
 	/// the total staked for rewards in the associated RewardEra
diff --git a/pallets/capacity/src/weights.rs b/pallets/capacity/src/weights.rs
index 4bb92e63ad..7bfce933f5 100644
--- a/pallets/capacity/src/weights.rs
+++ b/pallets/capacity/src/weights.rs
@@ -18,9 +18,9 @@
 //! Autogenerated weights for pallet_capacity
 //!
 //! THIS FILE WAS AUTO-GENERATED USING THE SUBSTRATE BENCHMARK CLI VERSION 4.0.0-dev
-//! DATE: 2024-03-06, STEPS: `20`, REPEAT: `10`, LOW RANGE: `[]`, HIGH RANGE: `[]`
+//! DATE: 2024-04-25, STEPS: `20`, REPEAT: `10`, LOW RANGE: `[]`, HIGH RANGE: `[]`
 //! WORST CASE MAP SIZE: `1000000`
-//! HOSTNAME: `benchmark-runner-wc6w8-pv9rm`, CPU: `Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz`
+//! HOSTNAME: `UL-Mac`, CPU: `<UNKNOWN>`
 //! EXECUTION: , WASM-EXECUTION: Compiled, CHAIN: Some("frequency-bench"), DB CACHE: 1024
 
 // Executed Command:
@@ -55,7 +55,6 @@ pub trait WeightInfo {
 	fn unstake() -> Weight;
 	fn set_epoch_length() -> Weight;
 	fn change_staking_target() -> Weight;
-
 	fn provider_boost() -> Weight;
 }
 
@@ -80,8 +79,8 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 		// Proof Size summary in bytes:
 		//  Measured:  `174`
 		//  Estimated: `6249`
-		// Minimum execution time: 39_050_000 picoseconds.
-		Weight::from_parts(39_836_000, 6249)
+		// Minimum execution time: 30_000_000 picoseconds.
+		Weight::from_parts(31_000_000, 6249)
 			.saturating_add(T::DbWeight::get().reads(7_u64))
 			.saturating_add(T::DbWeight::get().writes(4_u64))
 	}
@@ -117,6 +116,8 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 	}
 	/// Storage: `Capacity::StakingAccountLedger` (r:1 w:1)
 	/// Proof: `Capacity::StakingAccountLedger` (`max_values`: None, `max_size`: Some(57), added: 2532, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::StakingRewardPool` (r:1 w:1)
+	/// Proof: `Capacity::StakingRewardPool` (`max_values`: None, `max_size`: Some(60), added: 2535, mode: `MaxEncodedLen`)
 	/// Storage: `Capacity::UnstakeUnlocks` (r:1 w:1)
 	/// Proof: `Capacity::UnstakeUnlocks` (`max_values`: None, `max_size`: Some(121), added: 2596, mode: `MaxEncodedLen`)
 	/// Storage: `Capacity::StakingTargetLedger` (r:1 w:1)
@@ -125,12 +126,12 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 	/// Proof: `Capacity::CapacityLedger` (`max_values`: None, `max_size`: Some(68), added: 2543, mode: `MaxEncodedLen`)
 	fn unstake() -> Weight {
 		// Proof Size summary in bytes:
-		//  Measured:  `271`
+		//  Measured:  `343`
 		//  Estimated: `5071`
-		// Minimum execution time: 26_198_000 picoseconds.
-		Weight::from_parts(26_644_000, 5071)
-			.saturating_add(T::DbWeight::get().reads(4_u64))
-			.saturating_add(T::DbWeight::get().writes(4_u64))
+		// Minimum execution time: 23_000_000 picoseconds.
+		Weight::from_parts(24_000_000, 5071)
+			.saturating_add(T::DbWeight::get().reads(5_u64))
+			.saturating_add(T::DbWeight::get().writes(5_u64))
 	}
 	/// Storage: `Capacity::EpochLength` (r:0 w:1)
 	/// Proof: `Capacity::EpochLength` (`max_values`: Some(1), `max_size`: Some(4), added: 499, mode: `MaxEncodedLen`)
@@ -138,23 +139,53 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 		// Proof Size summary in bytes:
 		//  Measured:  `0`
 		//  Estimated: `0`
-		// Minimum execution time: 6_361_000 picoseconds.
-		Weight::from_parts(6_550_000, 0)
+		// Minimum execution time: 5_000_000 picoseconds.
+		Weight::from_parts(5_000_000, 0)
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
-
-	/// Storage:
-	/// Proof:
+	/// Storage: `Capacity::Retargets` (r:1 w:1)
+	/// Proof: `Capacity::Retargets` (`max_values`: None, `max_size`: Some(48), added: 2523, mode: `MaxEncodedLen`)
+	/// Storage: `Msa::ProviderToRegistryEntry` (r:1 w:0)
+	/// Proof: `Msa::ProviderToRegistryEntry` (`max_values`: None, `max_size`: Some(33), added: 2508, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::StakingAccountLedger` (r:1 w:0)
+	/// Proof: `Capacity::StakingAccountLedger` (`max_values`: None, `max_size`: Some(57), added: 2532, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::StakingTargetLedger` (r:2 w:2)
+	/// Proof: `Capacity::StakingTargetLedger` (`max_values`: None, `max_size`: Some(88), added: 2563, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::CapacityLedger` (r:2 w:2)
+	/// Proof: `Capacity::CapacityLedger` (`max_values`: None, `max_size`: Some(68), added: 2543, mode: `MaxEncodedLen`)
 	fn change_staking_target() -> Weight {
-		Weight::from_parts(1_000_000, 0)
-			.saturating_add(T::DbWeight::get().writes(1_u64))
+		// Proof Size summary in bytes:
+		//  Measured:  `315`
+		//  Estimated: `7601`
+		// Minimum execution time: 24_000_000 picoseconds.
+		Weight::from_parts(25_000_000, 7601)
+			.saturating_add(T::DbWeight::get().reads(7_u64))
+			.saturating_add(T::DbWeight::get().writes(5_u64))
 	}
-
-	/// Storage:
-	/// Proof:
+	/// Storage: `Msa::ProviderToRegistryEntry` (r:1 w:0)
+	/// Proof: `Msa::ProviderToRegistryEntry` (`max_values`: None, `max_size`: Some(33), added: 2508, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::StakingAccountLedger` (r:1 w:1)
+	/// Proof: `Capacity::StakingAccountLedger` (`max_values`: None, `max_size`: Some(57), added: 2532, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::StakingTargetLedger` (r:1 w:1)
+	/// Proof: `Capacity::StakingTargetLedger` (`max_values`: None, `max_size`: Some(88), added: 2563, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::CapacityLedger` (r:1 w:1)
+	/// Proof: `Capacity::CapacityLedger` (`max_values`: None, `max_size`: Some(68), added: 2543, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::StakingRewardPool` (r:1 w:1)
+	/// Proof: `Capacity::StakingRewardPool` (`max_values`: None, `max_size`: Some(60), added: 2535, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::UnstakeUnlocks` (r:1 w:0)
+	/// Proof: `Capacity::UnstakeUnlocks` (`max_values`: None, `max_size`: Some(121), added: 2596, mode: `MaxEncodedLen`)
+	/// Storage: `Balances::Freezes` (r:1 w:1)
+	/// Proof: `Balances::Freezes` (`max_values`: None, `max_size`: Some(85), added: 2560, mode: `MaxEncodedLen`)
+	/// Storage: `Balances::Locks` (r:1 w:0)
+	/// Proof: `Balances::Locks` (`max_values`: None, `max_size`: Some(1299), added: 3774, mode: `MaxEncodedLen`)
 	fn provider_boost() -> Weight {
-		Weight::from_parts(1_000_000, 0)
-			.saturating_add(T::DbWeight::get().writes(1_u64))
+		// Proof Size summary in bytes:
+		//  Measured:  `247`
+		//  Estimated: `6249`
+		// Minimum execution time: 34_000_000 picoseconds.
+		Weight::from_parts(35_000_000, 6249)
+			.saturating_add(T::DbWeight::get().reads(8_u64))
+			.saturating_add(T::DbWeight::get().writes(5_u64))
 	}
 }
 
@@ -215,6 +246,8 @@ impl WeightInfo for () {
 	}
 	/// Storage: `Capacity::StakingAccountLedger` (r:1 w:1)
 	/// Proof: `Capacity::StakingAccountLedger` (`max_values`: None, `max_size`: Some(57), added: 2532, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::StakingRewardPool` (r:1 w:1)
+	/// Proof: `Capacity::StakingRewardPool` (`max_values`: None, `max_size`: Some(60), added: 2535, mode: `MaxEncodedLen`)
 	/// Storage: `Capacity::UnstakeUnlocks` (r:1 w:1)
 	/// Proof: `Capacity::UnstakeUnlocks` (`max_values`: None, `max_size`: Some(121), added: 2596, mode: `MaxEncodedLen`)
 	/// Storage: `Capacity::StakingTargetLedger` (r:1 w:1)
@@ -223,12 +256,12 @@ impl WeightInfo for () {
 	/// Proof: `Capacity::CapacityLedger` (`max_values`: None, `max_size`: Some(68), added: 2543, mode: `MaxEncodedLen`)
 	fn unstake() -> Weight {
 		// Proof Size summary in bytes:
-		//  Measured:  `271`
+		//  Measured:  `343`
 		//  Estimated: `5071`
-		// Minimum execution time: 26_198_000 picoseconds.
-		Weight::from_parts(26_644_000, 5071)
-			.saturating_add(RocksDbWeight::get().reads(4_u64))
-			.saturating_add(RocksDbWeight::get().writes(4_u64))
+		// Minimum execution time: 23_000_000 picoseconds.
+		Weight::from_parts(24_000_000, 5071)
+			.saturating_add(RocksDbWeight::get().reads(5_u64))
+			.saturating_add(RocksDbWeight::get().writes(5_u64))
 	}
 	/// Storage: `Capacity::EpochLength` (r:0 w:1)
 	/// Proof: `Capacity::EpochLength` (`max_values`: Some(1), `max_size`: Some(4), added: 499, mode: `MaxEncodedLen`)
@@ -236,22 +269,52 @@ impl WeightInfo for () {
 		// Proof Size summary in bytes:
 		//  Measured:  `0`
 		//  Estimated: `0`
-		// Minimum execution time: 6_361_000 picoseconds.
-		Weight::from_parts(6_550_000, 0)
+		// Minimum execution time: 5_000_000 picoseconds.
+		Weight::from_parts(5_000_000, 0)
 			.saturating_add(RocksDbWeight::get().writes(1_u64))
 	}
-
-	/// Storage:
-	/// Proof:
+	/// Storage: `Capacity::Retargets` (r:1 w:1)
+	/// Proof: `Capacity::Retargets` (`max_values`: None, `max_size`: Some(48), added: 2523, mode: `MaxEncodedLen`)
+	/// Storage: `Msa::ProviderToRegistryEntry` (r:1 w:0)
+	/// Proof: `Msa::ProviderToRegistryEntry` (`max_values`: None, `max_size`: Some(33), added: 2508, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::StakingAccountLedger` (r:1 w:0)
+	/// Proof: `Capacity::StakingAccountLedger` (`max_values`: None, `max_size`: Some(57), added: 2532, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::StakingTargetLedger` (r:2 w:2)
+	/// Proof: `Capacity::StakingTargetLedger` (`max_values`: None, `max_size`: Some(88), added: 2563, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::CapacityLedger` (r:2 w:2)
+	/// Proof: `Capacity::CapacityLedger` (`max_values`: None, `max_size`: Some(68), added: 2543, mode: `MaxEncodedLen`)
 	fn change_staking_target() -> Weight {
-		Weight::from_parts(1_000_000, 0)
-			.saturating_add(RocksDbWeight::get().writes(1_u64))
+		// Proof Size summary in bytes:
+		//  Measured:  `315`
+		//  Estimated: `7601`
+		// Minimum execution time: 24_000_000 picoseconds.
+		Weight::from_parts(25_000_000, 7601)
+			.saturating_add(RocksDbWeight::get().reads(7_u64))
+			.saturating_add(RocksDbWeight::get().writes(5_u64))
 	}
-
-	/// Storage:
-	/// Proof:
+	/// Storage: `Msa::ProviderToRegistryEntry` (r:1 w:0)
+	/// Proof: `Msa::ProviderToRegistryEntry` (`max_values`: None, `max_size`: Some(33), added: 2508, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::StakingAccountLedger` (r:1 w:1)
+	/// Proof: `Capacity::StakingAccountLedger` (`max_values`: None, `max_size`: Some(57), added: 2532, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::StakingTargetLedger` (r:1 w:1)
+	/// Proof: `Capacity::StakingTargetLedger` (`max_values`: None, `max_size`: Some(88), added: 2563, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::CapacityLedger` (r:1 w:1)
+	/// Proof: `Capacity::CapacityLedger` (`max_values`: None, `max_size`: Some(68), added: 2543, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::StakingRewardPool` (r:1 w:1)
+	/// Proof: `Capacity::StakingRewardPool` (`max_values`: None, `max_size`: Some(60), added: 2535, mode: `MaxEncodedLen`)
+	/// Storage: `Capacity::UnstakeUnlocks` (r:1 w:0)
+	/// Proof: `Capacity::UnstakeUnlocks` (`max_values`: None, `max_size`: Some(121), added: 2596, mode: `MaxEncodedLen`)
+	/// Storage: `Balances::Freezes` (r:1 w:1)
+	/// Proof: `Balances::Freezes` (`max_values`: None, `max_size`: Some(85), added: 2560, mode: `MaxEncodedLen`)
+	/// Storage: `Balances::Locks` (r:1 w:0)
+	/// Proof: `Balances::Locks` (`max_values`: None, `max_size`: Some(1299), added: 3774, mode: `MaxEncodedLen`)
 	fn provider_boost() -> Weight {
-		Weight::from_parts(1_000_000, 0)
-			.saturating_add(RocksDbWeight::get().writes(1_u64))
+		// Proof Size summary in bytes:
+		//  Measured:  `247`
+		//  Estimated: `6249`
+		// Minimum execution time: 34_000_000 picoseconds.
+		Weight::from_parts(35_000_000, 6249)
+			.saturating_add(RocksDbWeight::get().reads(8_u64))
+			.saturating_add(RocksDbWeight::get().writes(5_u64))
 	}
 }