Improve hybrid staking unit tests

pallet/staking/src/lib.rs

@@ -789,6 +789,8 @@ pub mod pallet {
 			for n_exposure in c_exposure.nominators {
 				let n_payout = Perbill::from_rational(n_exposure.vote, c_exposure.vote) * n_payout;
 
+				dbg!(n_payout);
+
 				if collator == n_exposure.who {
 					// If the collator nominated themselves.
 
@@ -827,6 +829,19 @@ pub mod pallet {
 			let (n1, n2) = Self::elect_ns();
 			let cs_from_contract = Self::try_elect(n1, Self::elect_from_contract);
 			let cs_from_pallet = Self::try_elect(n2, Self::elect);
+
+			if n1 != cs_from_contract.len() as u32 || n2 != cs_from_pallet.len() as u32 {
+				log::error!(
+					"[pallet::staking] collator count mismatch; \
+					expected collator count from contract: {n1}, from pallet: {n2}, \
+					actual collator count from contract: {}, from pallet: {}",
+					cs_from_contract.len(),
+					cs_from_pallet.len(),
+				);
+
+				return None;
+			}
+
 			let cs = [cs_from_contract, cs_from_pallet].concat();
 
 			if cs.is_empty() {

pallet/staking/src/mock.rs

@@ -53,9 +53,14 @@ impl From<H160> for AccountId {
 		Self(H160::to_low_u64_le(&value))
 	}
 }
-impl Into<H160> for AccountId {
-	fn into(self) -> H160 {
-		H160::from_low_u64_le(self.0)
+impl From<AccountId> for u64 {
+	fn from(value: AccountId) -> Self {
+		value.0
+	}
+}
+impl From<AccountId> for H160 {
+	fn from(value: AccountId) -> Self {
+		H160::from_low_u64_le(value.0)
 	}
 }
 impl Display for AccountId {
@@ -257,10 +262,9 @@ impl crate::Stake for RingStaking {
 impl crate::Election<AccountId> for RingStaking {
 	fn elect(n: u32) -> Option<Vec<AccountId>> {
 		Some(
-			(100..=(100 + n) as u64)
+			(100..(100 + n) as u64)
 				.map(|i| {
 					let who = AccountId(i);
-					let _ = <pallet_balances::Pallet<Runtime>>::deposit_creating(&who, i as _);
 
 					assert_ok!(Session::set_keys(
 						RuntimeOrigin::signed(who),
@@ -412,8 +416,8 @@ impl Default for ExtBuilder {
 	}
 }
 
-pub fn preset_collators(n: u64) {
-	(10..(10 + n)).for_each(|i| {
+pub fn preset_collator_wait_list(n: u64) {
+	(10..10 + n).for_each(|i| {
 		let who = AccountId(i);
 		let _ = <pallet_balances::Pallet<Runtime>>::deposit_creating(&who, i as _);
 
@@ -426,6 +430,10 @@ pub fn preset_collators(n: u64) {
 		assert_ok!(Staking::collect(RuntimeOrigin::signed(who), Perbill::zero()));
 		assert_ok!(Staking::nominate(RuntimeOrigin::signed(who), who));
 	});
+	(100..100 + n).for_each(|i| {
+		let who = AccountId(i);
+		let _ = <pallet_balances::Pallet<Runtime>>::deposit_creating(&who, i as _);
+	});
 }
 
 pub fn initialize_block(number: BlockNumber) {

pallet/staking/src/tests.rs

@@ -797,7 +797,7 @@ fn elect_ns_should_work() {
 #[test]
 fn hybrid_election_should_work() {
 	ExtBuilder::default().collator_count(10).build().execute_with(|| {
-		mock::preset_collators(10);
+		mock::preset_collator_wait_list(10);
 		new_session();
 		new_session();
 
@@ -811,7 +811,7 @@ fn hybrid_election_should_work() {
 		new_session();
 
 		assert_eq!(
-			(100..103).map(AccountId).chain((12..20).rev().map(AccountId)).collect::<Vec<_>>(),
+			(100..102).map(AccountId).chain((12..20).rev().map(AccountId)).collect::<Vec<_>>(),
 			<pallet_session::Validators<Runtime>>::get()
 		);
 
@@ -820,7 +820,7 @@ fn hybrid_election_should_work() {
 		new_session();
 
 		assert_eq!(
-			(100..106).map(AccountId).chain((15..20).rev().map(AccountId)).collect::<Vec<_>>(),
+			(100..105).map(AccountId).chain((15..20).rev().map(AccountId)).collect::<Vec<_>>(),
 			<pallet_session::Validators<Runtime>>::get()
 		);
 
@@ -829,7 +829,7 @@ fn hybrid_election_should_work() {
 		new_session();
 
 		assert_eq!(
-			(100..108).map(AccountId).chain((17..20).rev().map(AccountId)).collect::<Vec<_>>(),
+			(100..107).map(AccountId).chain((17..20).rev().map(AccountId)).collect::<Vec<_>>(),
 			<pallet_session::Validators<Runtime>>::get()
 		);
 
@@ -838,18 +838,51 @@ fn hybrid_election_should_work() {
 		new_session();
 
 		assert_eq!(
-			(100..=110).map(AccountId).collect::<Vec<_>>(),
+			(100..110).map(AccountId).collect::<Vec<_>>(),
 			<pallet_session::Validators<Runtime>>::get()
 		);
 	});
 }
 
 #[test]
 fn hybrid_payout_should_work() {
-	ExtBuilder::default().collator_count(10).build().execute_with(|| {
-		mock::preset_collators(10);
+	ExtBuilder::default().collator_count(10).inflation_type(1).build().execute_with(|| {
+		mock::preset_collator_wait_list(10);
 		Timestamp::set_timestamp(30 * DAY_IN_MILLIS);
 		new_session();
 		new_session();
+
+		let collators =
+			(100..105).map(AccountId).chain((15..20).rev().map(AccountId)).collect::<Vec<_>>();
+		assert_eq!(collators, <pallet_session::Validators<Runtime>>::get());
+
+		let collators_balances = collators.iter().map(Balances::free_balance).collect::<Vec<_>>();
+		let session_duration = Duration::new(12 * 600, 0).as_millis();
+		Efflux::time(session_duration - <Period as Get<u64>>::get() as Moment);
+		Staking::note_authors(&collators);
+		new_session();
+		payout();
+
+		assert_eq!(
+			[
+				1_000 * UNIT,
+				1_000 * UNIT,
+				1_000 * UNIT,
+				1_000 * UNIT,
+				1_000 * UNIT,
+				1_000 * UNIT,
+				1_000 * UNIT,
+				1_000 * UNIT,
+				1_000 * UNIT,
+				1_000 * UNIT,
+			]
+			.as_slice(),
+			collators
+				.into_iter()
+				.map(Balances::free_balance)
+				.zip(collators_balances)
+				.map(|(a, b)| a - b)
+				.collect::<Vec<_>>()
+		);
 	});
 }