Try-state for Referenda pallet

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
 	);
 }

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<T::BlockNumber> for Pallet<T, I> {
+		#[cfg(feature = "try-runtime")]
+		fn try_state(_n: T::BlockNumber) -> 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,86 @@ 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(T::BlockNumber::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(())
+		})
+	}
 }

frame/referenda/src/migration.rs

@@ -188,6 +188,14 @@ pub mod test {
 		}
 	}
 
+	fn increment_referendum_index() {
+		ReferendumCount::<T, ()>::mutate(|x| {
+			let r = *x;
+			*x += 1;
+			r
+		});
+	}
+
 	#[test]
 	pub fn referendum_status_v0() {
 		// make sure the bytes of the encoded referendum v0 is decodable.
@@ -199,17 +207,19 @@ 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());
+			increment_referendum_index();
 			v0::ReferendumInfoFor::<T, ()>::insert(2, ongoing_v0);
 			// create and insert into the storage an approved referendum v0.
 			let approved_v0 = v0::ReferendumInfoOf::<T, ()>::Approved(
 				123,
 				Deposit { who: 1, amount: 10 },
 				Some(Deposit { who: 2, amount: 20 }),
 			);
+			increment_referendum_index();
 			v0::ReferendumInfoFor::<T, ()>::insert(5, approved_v0);
 			// run migration from v0 to v1.
 			v1::MigrateV0ToV1::<T, ()>::on_runtime_upgrade();

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)]