stake sets the staking type

Author: shannonwells

pallets/capacity/src/lib.rs

@@ -352,7 +352,8 @@ pub mod pallet {
 			let staker = ensure_signed(origin)?;
 
 			let (mut staking_account, actual_amount) =
-				Self::ensure_can_stake(&staker, target, amount, staking_type)?;
+				Self::ensure_can_stake(&staker, target, amount, &staking_type)?;
+			staking_account.staking_type = staking_type;
 
 			let capacity = Self::increase_stake_and_issue_capacity(
 				&staker,
@@ -462,7 +463,7 @@ impl<T: Config> Pallet<T> {
 		staker: &T::AccountId,
 		target: MessageSourceId,
 		amount: BalanceOf<T>,
-		staking_type: StakingType,
+		staking_type: &StakingType,
 	) -> Result<(StakingAccountDetails<T>, BalanceOf<T>), DispatchError> {
 		ensure!(amount > Zero::zero(), Error::<T>::ZeroAmountNotAllowed);
 		ensure!(T::TargetValidator::validate(target), Error::<T>::InvalidTarget);
@@ -472,7 +473,7 @@ impl<T: Config> Pallet<T> {
 		// if total > 0 the account exists already. Prevent the staking type from being changed.
 		if staking_account.total > Zero::zero() {
 			ensure!(
-				staking_account.staking_type == staking_type,
+				staking_account.staking_type == *staking_type,
 				Error::<T>::CannotChangeStakingType
 			);
 		}

pallets/capacity/src/tests/stake_and_deposit_tests.rs

@@ -1,15 +1,17 @@
 use super::{mock::*, testing_utils::*};
 use crate::{BalanceOf, CapacityDetails, Error, Event, StakingAccountDetails};
 use common_primitives::{
-	capacity::{Nontransferable, StakingType, StakingType::MaximumCapacity},
+	capacity::{
+		Nontransferable, StakingType,
+		StakingType::{MaximumCapacity, ProviderBoost},
+	},
 	msa::MessageSourceId,
 };
-use frame_benchmarking::AnalysisChoice::Max;
 use frame_support::{assert_noop, assert_ok, traits::WithdrawReasons};
 use sp_runtime::ArithmeticError;
 
 #[test]
-fn stake_works() {
+fn stake_max_capacity_works() {
 	new_test_ext().execute_with(|| {
 		let account = 200;
 		let target: MessageSourceId = 1;
@@ -53,6 +55,35 @@ fn stake_works() {
 	});
 }
 
+#[test]
+fn stake_rewards_works() {
+	new_test_ext().execute_with(|| {
+		let account = 200;
+		let target: MessageSourceId = 1;
+		let amount = 50;
+		let capacity = 5;
+		register_provider(target, String::from("Foo"));
+		assert_ok!(Capacity::stake(RuntimeOrigin::signed(account), target, amount, ProviderBoost));
+
+		// Check that StakingAccountLedger is updated.
+		let staking_account: StakingAccountDetails<Test> =
+			Capacity::get_staking_account_for(account).unwrap();
+
+		assert_eq!(staking_account.total, 50);
+		assert_eq!(staking_account.active, 50);
+		assert_eq!(staking_account.unlocking.len(), 0);
+		assert_eq!(staking_account.staking_type, ProviderBoost);
+		assert_eq!(staking_account.last_rewards_claimed_at, None);
+		assert_eq!(staking_account.stake_change_unlocking.len(), 0);
+
+		let events = staking_events();
+		assert_eq!(events.first().unwrap(), &Event::Staked { account, target, amount, capacity });
+
+		assert_eq!(Balances::locks(&account)[0].amount, amount);
+		assert_eq!(Balances::locks(&account)[0].reasons, WithdrawReasons::all().into());
+	});
+}
+
 #[test]
 fn stake_errors_invalid_target_when_target_is_not_registered_provider() {
 	new_test_ext().execute_with(|| {
@@ -331,7 +362,7 @@ fn ensure_can_stake_errors_with_zero_amount_not_allowed() {
 		let target: MessageSourceId = 1;
 		let amount = 0;
 		assert_noop!(
-			Capacity::ensure_can_stake(&account, target, amount, MaximumCapacity),
+			Capacity::ensure_can_stake(&account, target, amount, &MaximumCapacity),
 			Error::<Test>::ZeroAmountNotAllowed
 		);
 	});
@@ -369,7 +400,7 @@ fn ensure_can_stake_errors_invalid_target() {
 		let amount = 1;
 
 		assert_noop!(
-			Capacity::ensure_can_stake(&account, target, amount, MaximumCapacity),
+			Capacity::ensure_can_stake(&account, target, amount, &MaximumCapacity),
 			Error::<Test>::InvalidTarget
 		);
 	});
@@ -384,7 +415,7 @@ fn ensure_can_stake_errors_insufficient_staking_amount() {
 		register_provider(target, String::from("Foo"));
 
 		assert_noop!(
-			Capacity::ensure_can_stake(&account, target, amount, MaximumCapacity),
+			Capacity::ensure_can_stake(&account, target, amount, &MaximumCapacity),
 			Error::<Test>::InsufficientStakingAmount
 		);
 	});
@@ -400,7 +431,7 @@ fn ensure_can_stake_is_successful() {
 
 		let staking_details = StakingAccountDetails::<Test>::default();
 		assert_ok!(
-			Capacity::ensure_can_stake(&account, target, amount, MaximumCapacity),
+			Capacity::ensure_can_stake(&account, target, amount, &MaximumCapacity),
 			(staking_details, BalanceOf::<Test>::from(10u64))
 		);
 	});