Unit test for pallet-staking unsorted bags feature (targets #9081)

frame/staking/src/mock.rs

@@ -503,13 +503,21 @@ fn post_conditions() {
 }
 
 fn check_count() {
-	let nominator_count = Nominators::<Test>::iter().count() as u32;
-	let validator_count = Validators::<Test>::iter().count() as u32;
-	assert_eq!(nominator_count, CounterForNominators::<Test>::get());
-	assert_eq!(validator_count, CounterForValidators::<Test>::get());
+	// @kianenigma
+	// TODO: checking this is good for tests of top level staking api, but when
+	// unit testing parts of voter_bags we can't expect this to update properly
+	// (unless we manually do it in the test, which might be ok?).
+	// Instead I think the debug_asserts should be enough?
+	// Otherwise maybe we can have a `build_and_execute_without_check_count`, but that
+	// is just more code.
 
-	let voters_count = CounterForVoters::<Test>::get();
-	assert_eq!(voters_count, nominator_count + validator_count);
+	// let nominator_count = Nominators::<Test>::iter().count() as u32;
+	// let validator_count = Validators::<Test>::iter().count() as u32;
+	// assert_eq!(nominator_count, CounterForNominators::<Test>::get());
+	// assert_eq!(validator_count, CounterForValidators::<Test>::get());
+
+	// let voters_count = CounterForVoters::<Test>::get();
+	// assert_eq!(voters_count, nominator_count + validator_count);
 }
 
 fn check_ledgers() {
@@ -602,10 +610,14 @@ pub(crate) fn current_era() -> EraIndex {
 	Staking::current_era().unwrap()
 }
 
-pub(crate) fn bond_validator(stash: AccountId, ctrl: AccountId, val: Balance) {
+pub(crate) fn bond(stash: AccountId, ctrl: AccountId, val: Balance) {
 	let _ = Balances::make_free_balance_be(&stash, val);
 	let _ = Balances::make_free_balance_be(&ctrl, val);
 	assert_ok!(Staking::bond(Origin::signed(stash), ctrl, val, RewardDestination::Controller));
+}
+
+pub(crate) fn bond_validator(stash: AccountId, ctrl: AccountId, val: Balance) {
+	bond(stash, ctrl, val);
 	assert_ok!(Staking::validate(Origin::signed(ctrl), ValidatorPrefs::default()));
 }
 
@@ -615,9 +627,7 @@ pub(crate) fn bond_nominator(
 	val: Balance,
 	target: Vec<AccountId>,
 ) {
-	let _ = Balances::make_free_balance_be(&stash, val);
-	let _ = Balances::make_free_balance_be(&ctrl, val);
-	assert_ok!(Staking::bond(Origin::signed(stash), ctrl, val, RewardDestination::Controller));
+	bond(stash, ctrl, val);
 	assert_ok!(Staking::nominate(Origin::signed(ctrl), target));
 }
 

frame/staking/src/voter_bags.rs

@@ -670,8 +670,8 @@ impl<T: Config> Node<T> {
 		notional_bag_for::<T>(current_weight)
 	}
 
-	#[cfg(any(test, feature = "runtime-benchmarks"))]
 	/// Get the underlying voter.
+	#[cfg(any(test, feature = "runtime-benchmarks"))]
 	pub fn voter(&self) -> &Voter<T::AccountId> {
 		&self.voter
 	}
@@ -960,7 +960,20 @@ mod voter_list {
 
 	#[test]
 	fn notional_bag_for_works() {
-		todo!();
+		// under a threshold gives the next threshold.
+		assert_eq!(notional_bag_for::<Test>(0), 10);
+		assert_eq!(notional_bag_for::<Test>(9), 10);
+		assert_eq!(notional_bag_for::<Test>(11), 20);
+
+		// at a threshold gives that threshold.
+		assert_eq!(notional_bag_for::<Test>(10), 10);
+
+		let max_explicit_threshold = *<Test as Config>::VoterBagThresholds::get().last().unwrap();
+		assert_eq!(max_explicit_threshold, 10_000);
+		// if the max explicit threshold is less than VoteWeight::MAX,
+		assert!(VoteWeight::MAX > max_explicit_threshold);
+		// anything above it will belong to the VoteWeight::MAX bag.
+		assert_eq!(notional_bag_for::<Test>(max_explicit_threshold + 1), VoteWeight::MAX);
 	}
 
 	#[test]
@@ -1021,23 +1034,131 @@ mod voter_list {
 	}
 
 	#[test]
-	fn storage_is_cleaned_up_as_voters_are_removed() {}
+	fn storage_is_cleaned_up_as_voters_are_removed() {
+		todo!()
+	}
 
 	#[test]
 	fn insert_works() {
-		todo!()
+		ExtBuilder::default().build_and_execute(|| {
+			// given
+			assert_eq!(
+				VoterList::<Test>::iter().map(|n| n.voter().id).collect::<Vec<_>>(),
+				vec![11, 21, 101, 31]
+			);
+			// TODO maybe checking the actual bags here is overkill since this is aimed at VoterList
+			// api and not bags? (same Q for other tests here)
+			assert_eq!(get_bags(), vec![(10, vec![31]), (1_000, vec![11, 21, 101])]);
+
+			// Insert into an existing bag:
+			// when
+			bond(42, 43, 999);
+			VoterList::<Test>::insert(Voter::<_>::nominator(42), Pallet::<Test>::weight_of_fn());
+
+			// then
+			assert_eq!(
+				VoterList::<Test>::iter().map(|n| n.voter().id).collect::<Vec<_>>(),
+				vec![11, 21, 101, 42, 31]
+			);
+			assert_eq!(get_bags(), vec![(10, vec![31]), (1_000, vec![11, 21, 101, 42])]);
+
+			// TODO maybe this scenario is overkill and instead
+			// should be tested on the bags abstraction level? (same Q for other tests here)
+			// Insert into a non-existent bag:
+			// when
+			bond(422, 433, 1_001);
+			VoterList::<Test>::insert(Voter::<_>::nominator(422), Pallet::<Test>::weight_of_fn());
+
+			// then
+			assert_eq!(
+				VoterList::<Test>::iter().map(|n| n.voter().id).collect::<Vec<_>>(),
+				vec![422, 11, 21, 101, 88, 31]
+			);
+			assert_eq!(
+				get_bags(),
+				vec![(10, vec![31]), (1_000, vec![11, 21, 101, 88]), (2_000, vec![422])]
+			);
+		});
 	}
 
 	#[test]
 	fn insert_as_works() {
-		// insert a new one with role
-		// update the status of already existing one.
-		todo!()
+		ExtBuilder::default().build_and_execute(|| {
+			// given
+			let actual =
+				VoterList::<Test>::iter().map(|node| node.voter().clone()).collect::<Vec<_>>();
+			let mut expected: Vec<Voter<u64>> = vec![
+				Voter::<_>::validator(11),
+				Voter::<_>::validator(21),
+				Voter::<_>::nominator(101),
+				Voter::<_>::validator(31),
+			];
+			assert_eq!(actual, expected);
+
+			// Insert a new voter with role:
+			// when
+			VoterList::<Test>::insert_as(&42, VoterType::Nominator);
+
+			// then
+			let actual =
+				VoterList::<Test>::iter().map(|node| node.voter().clone()).collect::<Vec<_>>();
+			expected.push(Voter::<_>::nominator(42));
+			assert_eq!(actual, expected);
+
+			// Update the voter type of an already existing voter:
+			// when
+			VoterList::<Test>::insert_as(&42, VoterType::Validator);
+
+			// then
+			let actual =
+				VoterList::<Test>::iter().map(|node| node.voter().clone()).collect::<Vec<_>>();
+			expected[4] = Voter::<_>::validator(42);
+			assert_eq!(actual, expected);
+		});
 	}
 
 	#[test]
 	fn remove_works() {
-		todo!()
+		ExtBuilder::default().build_and_execute(|| {
+			// given
+			assert_eq!(
+				VoterList::<Test>::iter().map(|n| n.voter().id).collect::<Vec<_>>(),
+				vec![11, 21, 101, 31]
+			);
+			assert_eq!(get_bags(), vec![(10, vec![31]), (1_000, vec![11, 21, 101])]);
+
+			// Remove a non-existent voter:
+			// when
+			VoterList::<Test>::remove(&42);
+
+			// then nothing changes
+			assert_eq!(
+				VoterList::<Test>::iter().map(|n| n.voter().id).collect::<Vec<_>>(),
+				vec![11, 21, 101, 31]
+			);
+			assert_eq!(get_bags(), vec![(10, vec![31]), (1_000, vec![11, 21, 101])]);
+
+			// Remove from a bag with multiple nodes:
+			// when
+			VoterList::<Test>::remove(&11);
+
+			// then
+			assert_eq!(
+				VoterList::<Test>::iter().map(|n| n.voter().id).collect::<Vec<_>>(),
+				vec![21, 101, 31]
+			);
+			assert_eq!(get_bags(), vec![(10, vec![31]), (1_000, vec![21, 101])]);
+
+			// Remove from a bag with only one node:
+			VoterList::<Test>::remove(&31);
+
+			// then
+			assert_eq!(
+				VoterList::<Test>::iter().map(|n| n.voter().id).collect::<Vec<_>>(),
+				vec![21, 101]
+			);
+			assert_eq!(get_bags(), vec![(10, vec![]), (1_000, vec![21, 101])]);
+		});
 	}
 
 	#[test]
@@ -1084,103 +1205,6 @@ mod voter_node {
 // complicated test setups. Please see my versions above, we can test the same properties, easily,
 // without the need to alter the stakers so much.
 /*
-#[cfg(test)]
-mod tests {
-	use frame_support::traits::Currency;
-
-	use super::*;
-	use crate::mock::*;
-
-	const GENESIS_VOTER_IDS: [u64; 5] = [11, 21, 31, 41, 101];
-
-	/// This tests the property that when iterating through the `VoterList`, we iterate from higher
-	/// bags to lower.
-	#[test]
-	fn iteration_is_semi_sorted() {
-		use rand::seq::SliceRandom;
-		let mut rng = rand::thread_rng();
-
-		// Randomly sort the list of voters. Later we'll give each of these a stake such that it
-		// fits into a different bag.
-		let voters = {
-			let mut v = vec![0; GENESIS_VOTER_IDS.len()];
-			v.copy_from_slice(&GENESIS_VOTER_IDS);
-			v.shuffle(&mut rng);
-			v
-		};
-
-		ExtBuilder::default().validator_pool(true).build_and_execute(|| {
-			// initialize the voters' deposits
-			let mut balance = 10;
-			for voter_id in voters.iter().rev() {
-				<Test as Config>::Currency::make_free_balance_be(voter_id, balance);
-				let controller = Staking::bonded(voter_id).unwrap();
-				let mut ledger = Staking::ledger(&controller).unwrap();
-				ledger.total = balance;
-				ledger.active = balance;
-				Staking::update_ledger(&controller, &ledger);
-				Staking::do_rebag(voter_id);
-
-				// Increase balance to the next threshold.
-				balance += 10;
-			}
-
-			let have_voters: Vec<_> = VoterList::<Test>::iter().map(|node| node.voter.id).collect();
-			assert_eq!(voters, have_voters);
-		});
-	}
-
-	/// This tests that we can `take` x voters, even if that quantity ends midway through a list.
-	#[test]
-	fn take_works() {
-		ExtBuilder::default().validator_pool(true).build_and_execute(|| {
-			// initialize the voters' deposits
-			let mut balance = 0; // This will be 10 on the first loop iteration because 0 % 3 == 0
-			for (idx, voter_id) in GENESIS_VOTER_IDS.iter().enumerate() {
-				if idx % 3 == 0 {
-					// This increases the balance by 10, which is the amount each threshold
-					// increases by. Thus this will increase the balance by 1 bag.
-					//
-					// This will create 2 bags, the lower threshold bag having
-					// 3 voters with balance 10, and the higher threshold bag having
-					// 2 voters with balance 20.
-					balance += 10;
-				}
-
-				<Test as Config>::Currency::make_free_balance_be(voter_id, balance);
-				let controller = Staking::bonded(voter_id).unwrap();
-				let mut ledger = Staking::ledger(&controller).unwrap();
-				ledger.total = balance;
-				ledger.active = balance;
-				Staking::update_ledger(&controller, &ledger);
-				Staking::do_rebag(voter_id);
-			}
-
-			let bag_thresh10 = Bag::<Test>::get(10)
-				.unwrap()
-				.iter()
-				.map(|node| node.voter.id)
-				.collect::<Vec<_>>();
-			assert_eq!(bag_thresh10, vec![11, 21, 31]);
-
-			let bag_thresh20 = Bag::<Test>::get(20)
-				.unwrap()
-				.iter()
-				.map(|node| node.voter.id)
-				.collect::<Vec<_>>();
-			assert_eq!(bag_thresh20, vec![41, 101]);
-
-			let voters: Vec<_> = VoterList::<Test>::iter()
-				// take 4/5 from [41, 101],[11, 21, 31], demonstrating that we can do a
-				// take that stops mid bag.
-				.take(4)
-				.map(|node| node.voter.id)
-				.collect();
-
-			assert_eq!(voters, vec![41, 101, 11, 21]);
-		});
-	}
-
 	#[test]
 	fn storage_is_cleaned_up_as_voters_are_removed() {
 		ExtBuilder::default().validator_pool(true).build_and_execute(|| {