stake sets the staking type

pallets/capacity/src/lib.rs

@@ -350,7 +350,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,
@@ -460,7 +461,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);
@@ -470,7 +471,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))
  );
  });