Try-state for Referenda pallet (#13949)

* Try-state for Referenda pallet

* fix & more tests

* checking more stuff

* remove log

* two more tests

* Update frame/referenda/src/lib.rs

Co-authored-by: Muharem Ismailov <ismailov.m.h@gmail.com>

* fixes

* new check

* merge fixes

* fix warning

* remove check

* Update frame/referenda/src/lib.rs

Co-authored-by: Gonçalo Pestana <g6pestana@gmail.com>

* separate into multiple functions

* clean up

* unnecessary return value specified

* fix

* Update frame/referenda/src/lib.rs

* fmt

* remove import

* fmt

* fix CI

* Update frame/referenda/src/lib.rs

Co-authored-by: Liam Aharon <liam.aharon@hotmail.com>

* last fixes

* :P

* fmt

---------

Co-authored-by: Muharem Ismailov <ismailov.m.h@gmail.com>
Co-authored-by: Gonçalo Pestana <g6pestana@gmail.com>
Co-authored-by: Liam Aharon <liam.aharon@hotmail.com>

substrate/frame/referenda/src/benchmarking.rs

@@ -659,7 +659,7 @@ benchmarks_instance_pallet! {
 
 	impl_benchmark_test_suite!(
 		Referenda,
-		crate::mock::new_test_ext(),
+		crate::mock::ExtBuilder::default().build(),
 		crate::mock::Test
 	);
 }

substrate/frame/referenda/src/lib.rs

@@ -66,6 +66,7 @@
 
 use codec::{Codec, Encode};
 use frame_support::{
+	dispatch::DispatchResult,
 	ensure,
 	traits::{
 		schedule::{
@@ -416,6 +417,15 @@ pub mod pallet {
 		PreimageNotExist,
 	}
 
+	#[pallet::hooks]
+	impl<T: Config<I>, I: 'static> Hooks<BlockNumberFor<T>> for Pallet<T, I> {
+		#[cfg(feature = "try-runtime")]
+		fn try_state(_n: BlockNumberFor<T>) -> Result<(), sp_runtime::TryRuntimeError> {
+			Self::do_try_state()?;
+			Ok(())
+		}
+	}
+
 	#[pallet::call]
 	impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		/// Propose a referendum on a privileged action.
@@ -1032,7 +1042,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 	/// - If it's ready to be decided, start deciding;
 	/// - If it's not ready to be decided and non-deciding timeout has passed, fail;
 	/// - If it's ongoing and passing, ensure confirming; if at end of confirmation period, pass.
-	/// - If it's ongoing and not passing, stop confirning; if it has reached end time, fail.
+	/// - If it's ongoing and not passing, stop confirming; if it has reached end time, fail.
 	///
 	/// Weight will be a bit different depending on what it does, but it's designed so as not to
 	/// differ dramatically, especially if `MaxQueue` is kept small. In particular _there are no
@@ -1284,4 +1294,87 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 			Self::deposit_event(Event::<T, I>::MetadataCleared { index, hash });
 		}
 	}
+
+	/// Ensure the correctness of the state of this pallet.
+	///
+	/// The following assertions must always apply.
+	///
+	/// General assertions:
+	///
+	/// * [`ReferendumCount`] must always be equal to the number of referenda in
+	///   [`ReferendumInfoFor`].
+	/// * Referendum indices in [`MetadataOf`] must also be stored in [`ReferendumInfoFor`].
+	#[cfg(any(feature = "try-runtime", test))]
+	fn do_try_state() -> Result<(), sp_runtime::TryRuntimeError> {
+		ensure!(
+			ReferendumCount::<T, I>::get() as usize ==
+				ReferendumInfoFor::<T, I>::iter_keys().count(),
+			"Number of referenda in `ReferendumInfoFor` is different than `ReferendumCount`"
+		);
+
+		MetadataOf::<T, I>::iter_keys().try_for_each(|referendum_index| -> DispatchResult {
+			ensure!(
+				ReferendumInfoFor::<T, I>::contains_key(referendum_index),
+				"Referendum indices in `MetadataOf` must also be stored in `ReferendumInfoOf`"
+			);
+			Ok(())
+		})?;
+
+		Self::try_state_referenda_info()?;
+		Self::try_state_tracks()?;
+
+		Ok(())
+	}
+
+	/// Looking at referenda info:
+	///
+	/// - Data regarding ongoing phase:
+	///
+	/// * There must exist track info for the track of the referendum.
+	/// * The deciding stage has to begin before confirmation period.
+	/// * If alarm is set the nudge call has to be at most [`UndecidingTimeout`] blocks away
+	///  from the submission block.
+	#[cfg(any(feature = "try-runtime", test))]
+	fn try_state_referenda_info() -> Result<(), sp_runtime::TryRuntimeError> {
+		ReferendumInfoFor::<T, I>::iter().try_for_each(|(_, referendum)| {
+			match referendum {
+				ReferendumInfo::Ongoing(status) => {
+					ensure!(
+						Self::track(status.track).is_some(),
+						"No track info for the track of the referendum."
+					);
+
+					if let Some(deciding) = status.deciding {
+						ensure!(
+							deciding.since <
+								deciding.confirming.unwrap_or(BlockNumberFor::<T>::max_value()),
+							"Deciding status cannot begin before confirming stage."
+						)
+					}
+				},
+				_ => {},
+			}
+			Ok(())
+		})
+	}
+
+	/// Looking at tracks:
+	///
+	/// * The referendum indices stored in [`TrackQueue`] must exist as keys in the
+	///  [`ReferendumInfoFor`] storage map.
+	#[cfg(any(feature = "try-runtime", test))]
+	fn try_state_tracks() -> Result<(), sp_runtime::TryRuntimeError> {
+		T::Tracks::tracks().iter().try_for_each(|track| {
+			TrackQueue::<T, I>::get(track.0).iter().try_for_each(
+				|(referendum_index, _)| -> Result<(), sp_runtime::TryRuntimeError> {
+					ensure!(
+					ReferendumInfoFor::<T, I>::contains_key(referendum_index),
+					"`ReferendumIndex` inside the `TrackQueue` should be a key in `ReferendumInfoFor`"
+				);
+					Ok(())
+				},
+			)?;
+			Ok(())
+		})
+	}
 }

substrate/frame/referenda/src/migration.rs

@@ -199,10 +199,11 @@ pub mod test {
 
 	#[test]
 	fn migration_v0_to_v1_works() {
-		new_test_ext().execute_with(|| {
+		ExtBuilder::default().build_and_execute(|| {
 			// create and insert into the storage an ongoing referendum v0.
 			let status_v0 = create_status_v0();
 			let ongoing_v0 = v0::ReferendumInfoOf::<T, ()>::Ongoing(status_v0.clone());
+			ReferendumCount::<T, ()>::mutate(|x| x.saturating_inc());
 			v0::ReferendumInfoFor::<T, ()>::insert(2, ongoing_v0);
 			// create and insert into the storage an approved referendum v0.
 			let approved_v0 = v0::ReferendumInfoOf::<T, ()>::Approved(
@@ -210,6 +211,7 @@ pub mod test {
 				Deposit { who: 1, amount: 10 },
 				Some(Deposit { who: 2, amount: 20 }),
 			);
+			ReferendumCount::<T, ()>::mutate(|x| x.saturating_inc());
 			v0::ReferendumInfoFor::<T, ()>::insert(5, approved_v0);
 			// run migration from v0 to v1.
 			v1::MigrateV0ToV1::<T, ()>::on_runtime_upgrade();

substrate/frame/referenda/src/mock.rs

@@ -225,23 +225,32 @@ impl Config for Test {
 	type Tracks = TestTracksInfo;
 	type Preimages = Preimage;
 }
+pub struct ExtBuilder {}
 
-pub fn new_test_ext() -> sp_io::TestExternalities {
-	let mut t = frame_system::GenesisConfig::<Test>::default().build_storage().unwrap();
-	let balances = vec![(1, 100), (2, 100), (3, 100), (4, 100), (5, 100), (6, 100)];
-	pallet_balances::GenesisConfig::<Test> { balances }
-		.assimilate_storage(&mut t)
-		.unwrap();
-	let mut ext = sp_io::TestExternalities::new(t);
-	ext.execute_with(|| System::set_block_number(1));
-	ext
+impl Default for ExtBuilder {
+	fn default() -> Self {
+		Self {}
+	}
 }
 
-/// Execute the function two times, with `true` and with `false`.
-#[allow(dead_code)]
-pub fn new_test_ext_execute_with_cond(execute: impl FnOnce(bool) -> () + Clone) {
-	new_test_ext().execute_with(|| (execute.clone())(false));
-	new_test_ext().execute_with(|| execute(true));
+impl ExtBuilder {
+	pub fn build(self) -> sp_io::TestExternalities {
+		let mut t = frame_system::GenesisConfig::<Test>::default().build_storage().unwrap();
+		let balances = vec![(1, 100), (2, 100), (3, 100), (4, 100), (5, 100), (6, 100)];
+		pallet_balances::GenesisConfig::<Test> { balances }
+			.assimilate_storage(&mut t)
+			.unwrap();
+		let mut ext = sp_io::TestExternalities::new(t);
+		ext.execute_with(|| System::set_block_number(1));
+		ext
+	}
+
+	pub fn build_and_execute(self, test: impl FnOnce() -> ()) {
+		self.build().execute_with(|| {
+			test();
+			Referenda::do_try_state().unwrap();
+		})
+	}
 }
 
 #[derive(Encode, Debug, Decode, TypeInfo, Eq, PartialEq, Clone, MaxEncodedLen)]

substrate/frame/referenda/src/tests.rs

@@ -30,7 +30,7 @@ use pallet_balances::Error as BalancesError;
 
 #[test]
 fn params_should_work() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		assert_eq!(ReferendumCount::<Test>::get(), 0);
 		assert_eq!(Balances::free_balance(42), 0);
 		assert_eq!(Balances::total_issuance(), 600);
@@ -39,7 +39,7 @@ fn params_should_work() {
 
 #[test]
 fn basic_happy_path_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		// #1: submit
 		assert_ok!(Referenda::submit(
 			RuntimeOrigin::signed(1),
@@ -75,7 +75,7 @@ fn basic_happy_path_works() {
 
 #[test]
 fn insta_confirm_then_kill_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let r = Confirming { immediate: true }.create();
 		run_to(6);
 		assert_ok!(Referenda::kill(RuntimeOrigin::root(), r));
@@ -85,7 +85,7 @@ fn insta_confirm_then_kill_works() {
 
 #[test]
 fn confirm_then_reconfirm_with_elapsed_trigger_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let r = Confirming { immediate: false }.create();
 		assert_eq!(confirming_until(r), 8);
 		run_to(7);
@@ -99,7 +99,7 @@ fn confirm_then_reconfirm_with_elapsed_trigger_works() {
 
 #[test]
 fn instaconfirm_then_reconfirm_with_elapsed_trigger_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let r = Confirming { immediate: true }.create();
 		run_to(6);
 		assert_eq!(confirming_until(r), 7);
@@ -113,7 +113,7 @@ fn instaconfirm_then_reconfirm_with_elapsed_trigger_works() {
 
 #[test]
 fn instaconfirm_then_reconfirm_with_voting_trigger_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let r = Confirming { immediate: true }.create();
 		run_to(6);
 		assert_eq!(confirming_until(r), 7);
@@ -131,7 +131,7 @@ fn instaconfirm_then_reconfirm_with_voting_trigger_works() {
 
 #[test]
 fn voting_should_extend_for_late_confirmation() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let r = Passing.create();
 		run_to(10);
 		assert_eq!(confirming_until(r), 11);
@@ -142,7 +142,7 @@ fn voting_should_extend_for_late_confirmation() {
 
 #[test]
 fn should_instafail_during_extension_confirmation() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let r = Passing.create();
 		run_to(10);
 		assert_eq!(confirming_until(r), 11);
@@ -155,7 +155,7 @@ fn should_instafail_during_extension_confirmation() {
 
 #[test]
 fn confirming_then_fail_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let r = Failing.create();
 		// Normally ends at 5 + 4 (voting period) = 9.
 		assert_eq!(deciding_and_failing_since(r), 5);
@@ -170,7 +170,7 @@ fn confirming_then_fail_works() {
 
 #[test]
 fn queueing_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		// Submit a proposal into a track with a queue len of 1.
 		assert_ok!(Referenda::submit(
 			RuntimeOrigin::signed(5),
@@ -269,7 +269,7 @@ fn queueing_works() {
 
 #[test]
 fn alarm_interval_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let call =
 			<Test as Config>::Preimages::bound(CallOf::<Test, ()>::from(Call::nudge_referendum {
 				index: 0,
@@ -290,7 +290,7 @@ fn alarm_interval_works() {
 
 #[test]
 fn decision_time_is_correct() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let decision_time = |since: u64| {
 			Pallet::<Test>::decision_time(
 				&DecidingStatus { since: since.into(), confirming: None },
@@ -308,7 +308,7 @@ fn decision_time_is_correct() {
 
 #[test]
 fn auto_timeout_should_happen_with_nothing_but_submit() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		// #1: submit
 		assert_ok!(Referenda::submit(
 			RuntimeOrigin::signed(1),
@@ -329,7 +329,7 @@ fn auto_timeout_should_happen_with_nothing_but_submit() {
 
 #[test]
 fn tracks_are_distinguished() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		assert_ok!(Referenda::submit(
 			RuntimeOrigin::signed(1),
 			Box::new(RawOrigin::Root.into()),
@@ -390,7 +390,7 @@ fn tracks_are_distinguished() {
 
 #[test]
 fn submit_errors_work() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let h = set_balance_proposal_bounded(1);
 		// No track for Signed origins.
 		assert_noop!(
@@ -418,7 +418,7 @@ fn submit_errors_work() {
 
 #[test]
 fn decision_deposit_errors_work() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let e = Error::<Test>::NotOngoing;
 		assert_noop!(Referenda::place_decision_deposit(RuntimeOrigin::signed(2), 0), e);
 
@@ -440,7 +440,7 @@ fn decision_deposit_errors_work() {
 
 #[test]
 fn refund_deposit_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let e = Error::<Test>::BadReferendum;
 		assert_noop!(Referenda::refund_decision_deposit(RuntimeOrigin::signed(1), 0), e);
 
@@ -465,7 +465,7 @@ fn refund_deposit_works() {
 
 #[test]
 fn refund_submission_deposit_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		// refund of non existing referendum fails.
 		let e = Error::<Test>::BadReferendum;
 		assert_noop!(Referenda::refund_submission_deposit(RuntimeOrigin::signed(1), 0), e);
@@ -503,7 +503,7 @@ fn refund_submission_deposit_works() {
 
 #[test]
 fn cancel_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let h = set_balance_proposal_bounded(1);
 		assert_ok!(Referenda::submit(
 			RuntimeOrigin::signed(1),
@@ -522,7 +522,7 @@ fn cancel_works() {
 
 #[test]
 fn cancel_errors_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let h = set_balance_proposal_bounded(1);
 		assert_ok!(Referenda::submit(
 			RuntimeOrigin::signed(1),
@@ -540,7 +540,7 @@ fn cancel_errors_works() {
 
 #[test]
 fn kill_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let h = set_balance_proposal_bounded(1);
 		assert_ok!(Referenda::submit(
 			RuntimeOrigin::signed(1),
@@ -560,7 +560,7 @@ fn kill_works() {
 
 #[test]
 fn kill_errors_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let h = set_balance_proposal_bounded(1);
 		assert_ok!(Referenda::submit(
 			RuntimeOrigin::signed(1),
@@ -601,7 +601,7 @@ fn curve_handles_all_inputs() {
 
 #[test]
 fn set_metadata_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		use frame_support::traits::Hash as PreimageHash;
 		// invalid preimage hash.
 		let invalid_hash: PreimageHash = [1u8; 32].into();
@@ -649,7 +649,7 @@ fn set_metadata_works() {
 
 #[test]
 fn clear_metadata_works() {
-	new_test_ext().execute_with(|| {
+	ExtBuilder::default().build_and_execute(|| {
 		let hash = note_preimage(1);
 		assert_ok!(Referenda::submit(
 			RuntimeOrigin::signed(1),