mega cleanup of staking tests

frame/staking/src/mock.rs

@@ -285,33 +285,35 @@ pub(crate) type StakingCall = crate::Call<Test>;
 pub(crate) type TestRuntimeCall = <Test as frame_system::Config>::Call;
 
 pub struct ExtBuilder {
-	validator_pool: bool,
 	nominate: bool,
 	validator_count: u32,
 	minimum_validator_count: u32,
-	fair: bool,
-	num_validators: Option<u32>,
 	invulnerables: Vec<AccountId>,
 	has_stakers: bool,
 	initialize_first_session: bool,
 	min_nominator_bond: Balance,
 	min_validator_bond: Balance,
+	balance_factor: Balance,
+	status: BTreeMap<AccountId, StakerStatus<AccountId>>,
+	stakes: BTreeMap<AccountId, Balance>,
+	stakers: Vec<(AccountId, AccountId, Balance, StakerStatus<AccountId>)>,
 }
 
 impl Default for ExtBuilder {
 	fn default() -> Self {
 		Self {
-			validator_pool: false,
 			nominate: true,
 			validator_count: 2,
 			minimum_validator_count: 0,
-			fair: true,
-			num_validators: None,
+			balance_factor: 1,
 			invulnerables: vec![],
 			has_stakers: true,
 			initialize_first_session: true,
 			min_nominator_bond: ExistentialDeposit::get(),
 			min_validator_bond: ExistentialDeposit::get(),
+			status: Default::default(),
+			stakes: Default::default(),
+			stakers: Default::default(),
 		}
 	}
 }
@@ -321,10 +323,6 @@ impl ExtBuilder {
 		EXISTENTIAL_DEPOSIT.with(|v| *v.borrow_mut() = existential_deposit);
 		self
 	}
-	pub fn validator_pool(mut self, validator_pool: bool) -> Self {
-		self.validator_pool = validator_pool;
-		self
-	}
 	pub fn nominate(mut self, nominate: bool) -> Self {
 		self.nominate = nominate;
 		self
@@ -341,14 +339,6 @@ impl ExtBuilder {
 		SLASH_DEFER_DURATION.with(|v| *v.borrow_mut() = eras);
 		self
 	}
-	pub fn fair(mut self, is_fair: bool) -> Self {
-		self.fair = is_fair;
-		self
-	}
-	pub fn num_validators(mut self, num_validators: u32) -> Self {
-		self.num_validators = Some(num_validators);
-		self
-	}
 	pub fn invulnerables(mut self, invulnerables: Vec<AccountId>) -> Self {
 		self.invulnerables = invulnerables;
 		self
@@ -381,41 +371,60 @@ impl ExtBuilder {
 		self.min_validator_bond = amount;
 		self
 	}
+	pub fn set_status(mut self, who: AccountId, status: StakerStatus<AccountId>) -> Self {
+		self.status.insert(who, status);
+		self
+	}
+	pub fn set_stake(mut self, who: AccountId, stake: Balance) -> Self {
+		self.stakes.insert(who, stake);
+		self
+	}
+	pub fn add_staker(
+		mut self,
+		stash: AccountId,
+		ctrl: AccountId,
+		stake: Balance,
+		status: StakerStatus<AccountId>,
+	) -> Self {
+		self.stakers.push((stash, ctrl, stake, status));
+		self
+	}
+	pub fn balance_factor(mut self, factor: Balance) -> Self {
+		self.balance_factor = factor;
+		self
+	}
 	fn build(self) -> sp_io::TestExternalities {
 		sp_tracing::try_init_simple();
 		let mut storage = frame_system::GenesisConfig::default().build_storage::<Test>().unwrap();
-		let balance_factor = if ExistentialDeposit::get() > 1 { 256 } else { 1 };
-
-		let num_validators = self.num_validators.unwrap_or(self.validator_count);
-		// Check that the number of validators is sensible.
-		assert!(num_validators <= 8);
-		let validators =
-			(0..num_validators).map(|x| ((x + 1) * 10 + 1) as AccountId).collect::<Vec<_>>();
 
 		let _ = pallet_balances::GenesisConfig::<Test> {
 			balances: vec![
-				(1, 10 * balance_factor),
-				(2, 20 * balance_factor),
-				(3, 300 * balance_factor),
-				(4, 400 * balance_factor),
-				(10, balance_factor),
-				(11, balance_factor * 1000),
-				(20, balance_factor),
-				(21, balance_factor * 2000),
-				(30, balance_factor),
-				(31, balance_factor * 2000),
-				(40, balance_factor),
-				(41, balance_factor * 2000),
-				(50, balance_factor),
-				(51, balance_factor * 2000),
-				(60, balance_factor),
-				(61, balance_factor * 2000),
-				(70, balance_factor),
-				(71, balance_factor * 2000),
-				(80, balance_factor),
-				(81, balance_factor * 2000),
-				(100, 2000 * balance_factor),
-				(101, 2000 * balance_factor),
+				(1, 10 * self.balance_factor),
+				(2, 20 * self.balance_factor),
+				(3, 300 * self.balance_factor),
+				(4, 400 * self.balance_factor),
+				// controllers
+				(10, self.balance_factor),
+				(20, self.balance_factor),
+				(30, self.balance_factor),
+				(40, self.balance_factor),
+				(50, self.balance_factor),
+				// stashes
+				(11, self.balance_factor * 1000),
+				(21, self.balance_factor * 2000),
+				(31, self.balance_factor * 2000),
+				(41, self.balance_factor * 2000),
+				(51, self.balance_factor * 2000),
+				// optional nominator
+				(100, self.balance_factor * 2000),
+				(101, self.balance_factor * 2000),
+				// aux accounts
+				(60, self.balance_factor),
+				(61, self.balance_factor * 2000),
+				(70, self.balance_factor),
+				(71, self.balance_factor * 2000),
+				(80, self.balance_factor),
+				(81, self.balance_factor * 2000),
 				// This allows us to have a total_payout different from 0.
 				(999, 1_000_000_000_000),
 			],
@@ -424,24 +433,45 @@ impl ExtBuilder {
 
 		let mut stakers = vec![];
 		if self.has_stakers {
-			let stake_21 = if self.fair { 1000 } else { 2000 };
-			let stake_31 = if self.validator_pool { balance_factor * 1000 } else { 1 };
-			let status_41 = if self.validator_pool {
-				StakerStatus::<AccountId>::Validator
-			} else {
-				StakerStatus::<AccountId>::Idle
-			};
-			let nominated = if self.nominate { vec![11, 21] } else { vec![] };
 			stakers = vec![
-				// (stash, controller, staked_amount, status)
-				(11, 10, balance_factor * 1000, StakerStatus::<AccountId>::Validator),
-				(21, 20, stake_21, StakerStatus::<AccountId>::Validator),
-				(31, 30, stake_31, StakerStatus::<AccountId>::Validator),
-				(41, 40, balance_factor * 1000, status_41),
-				// nominator
-				(101, 100, balance_factor * 500, StakerStatus::<AccountId>::Nominator(nominated)),
+				// (stash, ctrl, stake, status)
+				// these two will be elected in the default test where we elect 2.
+				(11, 10, self.balance_factor * 1000, StakerStatus::<AccountId>::Validator),
+				(21, 20, self.balance_factor * 1000, StakerStatus::<AccountId>::Validator),
+				// a loser validator
+				(31, 30, self.balance_factor * 500, StakerStatus::<AccountId>::Validator),
+				// an idle validator
+				(41, 40, self.balance_factor * 1000, StakerStatus::<AccountId>::Idle),
 			];
+			// optionally add a nominator
+			if self.nominate {
+				stakers.push((
+					101,
+					100,
+					self.balance_factor * 500,
+					StakerStatus::<AccountId>::Nominator(vec![11, 21]),
+				))
+			}
+			// replace any of the status if needed.
+			self.status.into_iter().for_each(|(stash, status)| {
+				if let Some((_, _, _, ref mut prev_status)) =
+					stakers.iter_mut().find(|s| s.0 == stash)
+				{
+					*prev_status = status
+				}
+			});
+			// replaced any of the stakes if needed.
+			self.stakes.into_iter().for_each(|(stash, stake)| {
+				if let Some((_, _, ref mut prev_stake, _)) =
+					stakers.iter_mut().find(|s| s.0 == stash)
+				{
+					*prev_stake = stake
+				}
+			});
+			// extend stakers if needed.
+			stakers.extend(self.stakers)
 		}
+
 		let _ = staking::GenesisConfig::<Test> {
 			stakers,
 			validator_count: self.validator_count,
@@ -455,10 +485,15 @@ impl ExtBuilder {
 		.assimilate_storage(&mut storage);
 
 		let _ = pallet_session::GenesisConfig::<Test> {
-			keys: validators
-				.iter()
-				.map(|x| (*x, *x, SessionKeys { other: UintAuthorityId(*x as u64) }))
-				.collect(),
+			keys: if self.has_stakers {
+				// genesis election will overwrite this, no worries.
+				Default::default()
+			} else {
+				// set some dummy validators in genesis.
+				(0..self.validator_count as u64)
+					.map(|x| (x, x, SessionKeys { other: UintAuthorityId(x as u64) }))
+					.collect()
+			},
 		}
 		.assimilate_storage(&mut storage);
 

frame/staking/src/pallet/mod.rs

@@ -497,13 +497,13 @@ pub mod pallet {
 					T::Currency::free_balance(&stash) >= balance,
 					"Stash does not have enough balance to bond."
 				);
-				let _ = <Pallet<T>>::bond(
+				frame_support::assert_ok!(<Pallet<T>>::bond(
 					T::Origin::from(Some(stash.clone()).into()),
 					T::Lookup::unlookup(controller.clone()),
 					balance,
 					RewardDestination::Staked,
-				);
-				let _ = match status {
+				));
+				frame_support::assert_ok!(match status {
 					StakerStatus::Validator => <Pallet<T>>::validate(
 						T::Origin::from(Some(controller.clone()).into()),
 						Default::default(),
@@ -513,7 +513,7 @@ pub mod pallet {
 						votes.iter().map(|l| T::Lookup::unlookup(l.clone())).collect(),
 					),
 					_ => Ok(()),
-				};
+				});
 			}
 		}
 	}