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

frame/staking/src/lib.rs

@@ -766,7 +766,7 @@ pub mod migrations {
 		pub fn pre_migrate<T: Config>() -> Result<(), &'static str> {
 			ensure!(StorageVersion::<T>::get() == Releases::V7_0_0, "must upgrade linearly");
 			ensure!(
-				VoterList::<T>::decode_len().unwrap_or_default() == 0,
+				VoterList::<T>::iter().count() == 0,
 				"voter list already exists"
 			);
 			Ok(())
@@ -1781,6 +1781,7 @@ pub mod pallet {
 				let era = Self::current_era().unwrap_or(0) + T::BondingDuration::get();
 				ledger.unlocking.push(UnlockChunk { value, era });
 				Self::update_ledger(&controller, &ledger);
+				Self::do_rebag(&ledger.stash);
 				Self::deposit_event(Event::<T>::Unbonded(ledger.stash, value));
 			}
 			Ok(())
@@ -2273,6 +2274,7 @@ pub mod pallet {
 
 			Self::deposit_event(Event::<T>::Bonded(ledger.stash.clone(), value));
 			Self::update_ledger(&controller, &ledger);
+			Self::do_rebag(&ledger.stash);
 			Ok(Some(
 				35 * WEIGHT_PER_MICROS +
 					50 * WEIGHT_PER_NANOS * (ledger.unlocking.len() as Weight) +

frame/staking/src/mock.rs

@@ -493,13 +493,23 @@ impl ExtBuilder {
 		ext.execute_with(test);
 		ext.execute_with(post_conditions);
 	}
+	/// WARNING: This should only be use for testing `VoterList` api or lower.
+	pub fn build_and_execute_without_check_count(self, test: impl FnOnce() -> ()) {
+		let mut ext = self.build();
+		ext.execute_with(test);
+		ext.execute_with(post_conditions_without_check_count);
+	}
 }
 
 fn post_conditions() {
+	post_conditions_without_check_count();
+	check_count();
+}
+
+fn post_conditions_without_check_count() {
 	check_nominators();
 	check_exposures();
 	check_ledgers();
-	check_count();
 }
 
 fn check_count() {
@@ -602,10 +612,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 +629,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));
 }
 
@@ -833,3 +845,26 @@ pub(crate) fn get_bags() -> Vec<(VoteWeight, Vec<AccountId>)> {
 		})
 		.collect::<Vec<_>>()
 }
+
+pub(crate) fn bag_as_ids(bag: &Bag<Test>) -> Vec<AccountId> {
+	bag.iter().map(|n| n.voter().id).collect::<Vec<_>>()
+}
+
+pub(crate) fn get_voter_list_as_ids() -> Vec<AccountId> {
+	VoterList::<Test>::iter().map(|n| n.voter().id).collect::<Vec<_>>()
+}
+
+pub(crate) fn get_voter_list_as_voters() -> Vec<voter_bags::Voter<AccountId>> {
+	VoterList::<Test>::iter().map(|node| node.voter().clone()).collect::<Vec<_>>()
+}
+
+// Useful for when you want to change the effectively bonded value but you don't want to use
+// the bond extrinsics because they implicitly rebag.
+pub(crate) fn set_ledger_and_free_balance(account: &AccountId, value: Balance) {
+	Balances::make_free_balance_be(account, value);
+	let controller = Staking::bonded(account).unwrap();
+	let mut ledger = Staking::ledger(&controller).unwrap();
+	ledger.total = value;
+	ledger.active = value;
+	Staking::update_ledger(&controller, &ledger);
+}

frame/staking/src/tests.rs

@@ -3827,88 +3827,87 @@ fn on_finalize_weight_is_nonzero() {
 
 // end-to-end nodes of the voter bags operation.
 mod voter_bags {
+	use super::Origin;
+	use crate::{mock::*, ValidatorPrefs};
+	use frame_support::{assert_ok, traits::Currency};
 
 	#[test]
-	fn rebag_works() {
-		todo!()
-	}
-}
-/*
-// TODO: this needs some love, retire it in favour of the one above. Use the mock data, don't make
-// it complicated with data setup, use the simplest data possible, instead check multiple
-// edge-cases.
-#[test]
-fn test_rebag() {
-	use crate::{
-		testing_utils::create_stash_controller,
-		voter_bags::{Bag, Node},
-	};
-	use frame_system::RawOrigin;
-
-	/// Make a validator and return its stash
-	fn make_validator(n: u32, balance_factor: u32) -> Result<AccountIdOf<Test>, &'static str> {
-		let (stash, controller) = create_stash_controller::<Test>(n, balance_factor, Default::default()).unwrap();
-
-		// Bond the full value of the stash
-		//
-		// By default, `create_stash_controller` only bonds 10% of the stash. However, we're going
-		// to want to edit one account's bonded value to match another's, so it's simpler if 100% of
-		// the balance is bonded.
-		let balance = <Test as Config>::Currency::free_balance(&stash);
-		Staking::bond_extra(RawOrigin::Signed(stash.clone()).into(), balance).unwrap();
-		Staking::validate(
-			RawOrigin::Signed(controller.clone()).into(),
-			ValidatorPrefs::default(),
-		).unwrap();
-
-		Ok(stash)
+	fn insert_and_remove_works() {
+		// we test insert/remove indirectly via `validate`, `nominate`, and chill
+		ExtBuilder::default().build_and_execute(|| {
+			// given
+			assert_eq!(get_bags(), vec![(10, vec![31]), (1000, vec![11, 21, 101])]);
+
+			// `bond`
+			bond(42, 43, 2_000);
+			// does not insert the voter
+			assert_eq!(get_bags(), vec![(10, vec![31]), (1000, vec![11, 21, 101])]);
+
+			// `validate`
+			assert_ok!(Staking::validate(Origin::signed(43).into(), ValidatorPrefs::default()));
+			// moves the voter into a bag
+			assert_eq!(
+				get_bags(),
+				vec![(10, vec![31]), (1000, vec![11, 21, 101]), (2000, vec![42])]
+			);
+
+			// `nominate`-ing, but not changing active stake (which implicitly calls remove)
+			assert_ok!(Staking::nominate(Origin::signed(43), vec![11]));
+			// does not change the voters position
+			assert_eq!(
+				get_bags(),
+				vec![(10, vec![31]), (1000, vec![11, 21, 101]), (2000, vec![42])]
+			);
+
+			// `chill`
+			assert_ok!(Staking::chill(Origin::signed(43)));
+			// removes the voter
+			assert_eq!(get_bags(), vec![(10, vec![31]), (1000, vec![11, 21, 101])]);
+		});
 	}
 
-	ExtBuilder::default().build_and_execute(|| {
-		// We want to have two validators: one, `stash`, is the one we will rebag.
-		// The other, `other_stash`, exists only so that the destination bag is not empty.
-		let stash = make_validator(0, 2000).unwrap();
-		let other_stash = make_validator(1, 9000).unwrap();
-
-		// verify preconditions
-		let weight_of = Staking::weight_of_fn();
-		let node = Node::<Test>::from_id(&stash).unwrap();
-		assert_eq!(
-			{
-				let origin_bag = Bag::<Test>::get(node.bag_upper).unwrap();
-				origin_bag.iter().count()
-			},
-			1,
-			"stash should be the only node in origin bag",
-		);
-		let other_node = Node::<Test>::from_id(&other_stash).unwrap();
-		assert!(!other_node.is_misplaced(&weight_of), "other stash balance never changed");
-		assert_ne!(
-			{
-				let destination_bag = Bag::<Test>::get(other_node.bag_upper);
-				destination_bag.iter().count()
-			},
-			0,
-			"destination bag should not be empty",
-		);
+	#[test]
+	fn rebag_works() {
+		ExtBuilder::default().build_and_execute(|| {
+			// add a nominator to genesis state
+			bond_nominator(42, 43, 20, vec![11]);
+			Balances::make_free_balance_be(&42, 2_000);
 
-		// Update `stash`'s value to match `other_stash`, and bond extra to update its weight.
-		//
-		// This implicitly calls rebag, so the user stays in the best bag they qualify for.
-		let new_balance = <Test as Config>::Currency::free_balance(&other_stash);
-		<Test as Config>::Currency::make_free_balance_be(&stash, new_balance);
-		Staking::bond_extra(
-			RawOrigin::Signed(stash.clone()).into(),
-			new_balance,
-		).unwrap();
-
-		// node should no longer be misplaced
-		// note that we refresh the node, in case the storage value has changed
-		let node = Node::<Test>::from_id(&stash).unwrap();
-		assert!(!node.is_misplaced(&weight_of), "node must be in proper place after rebag");
-	});
+			// given
+			assert_eq!(get_bags(), vec![(10, vec![31]), (20, vec![42]), (1000, vec![11, 21, 101])]);
+
+			// increase stake and implicitly rebag with `bond_extra` to the level of non-existent bag
+			assert_ok!(Staking::bond_extra(Origin::signed(42), 1_980)); // 20 + 1_980 = 2_000
+			assert_eq!(
+				get_bags(),
+				vec![(10, vec![31]), (1000, vec![11, 21, 101]), (2000, vec![42])]
+			);
+
+			// decrease stake within the range of the current bag
+			assert_ok!(Staking::unbond(Origin::signed(43), 999)); // 2000 - 999 = 1001
+			// does not change bags
+			assert_eq!(
+				get_bags(),
+				vec![(10, vec![31]), (1000, vec![11, 21, 101]), (2000, vec![42])]
+			);
+
+			// reduce stake to the level of a non-existent bag
+			assert_ok!(Staking::unbond(Origin::signed(43), 971)); // 1001 - 971 = 30
+			// creates the bag and moves the voter into it
+			assert_eq!(
+				get_bags(),
+				vec![(10, vec![31]), (30, vec![42]), (1000, vec![11, 21, 101]),]
+			);
+
+			// increase stake by `rebond`-ing to the level of a pre-existing bag
+			assert_ok!(Staking::rebond(Origin::signed(43), 31)); // 30 + 41 = 61
+			// moves the voter to that bag
+			assert_eq!(get_bags(), vec![(10, vec![31]), (1000, vec![11, 21, 101, 42]),]);
+
+			// TODO test rebag directly
+		});
+	}
 }
-*/
 
 mod election_data_provider {
 	use super::*;