[AHM] fix scheduler and checking balance tests

integration-tests/ahm/src/balances_test.rs

@@ -79,12 +79,15 @@ impl AhMigrationCheck for BalancesCrossChecker {
 			&<AhRuntime as pallet_ah_migrator::Config>::CheckingAccount::get(),
 		);
 
-		// AH checking account has incorrect 0.01 DOT balance because of the DED airdrop which
-		// added DOT ED to all existing AH accounts.
+		// Polkadot AH checking account has incorrect 0.01 DOT balance because of the DED airdrop
+		// which added DOT ED to all existing AH accounts.
 		// This is fine, we can just ignore/accept this small amount.
+		#[cfg(feature = "polkadot-ahm")]
 		defensive_assert!(
 			ah_checking_balance_before == pallet_balances::Pallet::<AhRuntime>::minimum_balance()
 		);
+		#[cfg(feature = "kusama-ahm")]
+		defensive_assert!(ah_checking_balance_before == 0);
 
 		(ah_total_issuance_before, ah_checking_balance_before)
 	}

pallets/ah-migrator/src/account.rs

@@ -285,7 +285,10 @@ pub mod tests {
 			// AH checking account has incorrect 0.01 DOT balance because of the DED airdrop which
 			// added DOT ED to all existing AH accounts.
 			// This is fine, we can just ignore/accept this small amount.
+			#[cfg(feature = "polkadot-ahm")]
 			defensive_assert!(checking_balance == <T as Config>::Currency::minimum_balance());
+			#[cfg(feature = "kusama-ahm")]
+			defensive_assert!(checking_balance.is_zero());
 
 			let mut ah_pre_payload = BTreeMap::new();
 			for (account, _) in frame_system::Account::<T>::iter() {

pallets/ah-migrator/src/benchmarking.rs

@@ -82,7 +82,7 @@ pub mod benchmarks {
 			);
 			let _ = <<T as pallet_multisig::Config>::Currency>::reserve(&creator, deposit).unwrap();
 
-			RcMultisig { creator, deposit, details: Some([2u8; 32].into()) }
+			RcMultisig { creator, deposit }
 		};
 
 		let messages = (0..n).map(|i| create_multisig(i.try_into().unwrap())).collect::<Vec<_>>();

pallets/ah-migrator/src/lib.rs

@@ -63,8 +63,8 @@ pub mod xcm_translation;
 pub use pallet::*;
 pub use pallet_rc_migrator::{
 	types::{
-		BenchmarkingDefault, ExceptResponseFor, LeftOrRight, MaxOnIdleOrInner,
-		QueuePriority as DmpQueuePriority, RouteInnerWithException,
+		BenchmarkingDefault, ExceptResponseFor, LeftIfFinished, LeftIfPending, LeftOrRight,
+		MaxOnIdleOrInner, QueuePriority as DmpQueuePriority, RouteInnerWithException,
 	},
 	weights_ah,
 };

pallets/ah-migrator/src/scheduler.rs

@@ -156,14 +156,14 @@ pub struct RcPrePayload<T: Config> {
 #[cfg(feature = "std")]
 impl<T: Config> crate::types::AhMigrationCheck for SchedulerMigrator<T> {
 	type RcPrePayload = Vec<u8>;
-	type AhPrePayload = ();
+	type AhPrePayload = Option<BlockNumberFor<T>>;
 
 	fn pre_check(_rc_pre_payload: Self::RcPrePayload) -> Self::AhPrePayload {
 		// Assert storage 'Scheduler::IncompleteSince::ah_pre::empty'
-		assert!(
-			pallet_scheduler::IncompleteSince::<T>::get().is_none(),
-			"IncompleteSince should be empty on asset hub before migration"
-		);
+
+		// since the scheduler pallet will run on Asset Hub before migration, it will set up its
+		// `IncompleteSince` storage value with `on_initialize` hook.
+		let incomplete_since = pallet_scheduler::IncompleteSince::<T>::get();
 
 		// Assert storage 'Scheduler::Agenda::ah_pre::empty'
 		assert!(
@@ -182,18 +182,28 @@ impl<T: Config> crate::types::AhMigrationCheck for SchedulerMigrator<T> {
 			pallet_scheduler::Retries::<T>::iter().next().is_none(),
 			"Retries map should be empty on asset hub before migration"
 		);
+
+		incomplete_since
 	}
 
-	fn post_check(rc_pre_payload: Self::RcPrePayload, _ah_pre_payload: Self::AhPrePayload) {
+	fn post_check(rc_pre_payload: Self::RcPrePayload, ah_incomplete_since: Self::AhPrePayload) {
 		let rc_payload = RcPrePayload::<T>::decode(&mut &rc_pre_payload[..])
 			.expect("Failed to decode RcPrePayload bytes");
 
 		// Assert storage 'Scheduler::IncompleteSince::ah_post::correct'
-		assert_eq!(
-			pallet_scheduler::IncompleteSince::<T>::get(),
-			rc_payload.incomplete_since,
-			"IncompleteSince on Asset Hub should match the RC value"
-		);
+		if rc_payload.incomplete_since.is_some() {
+			assert_eq!(
+				pallet_scheduler::IncompleteSince::<T>::get(),
+				rc_payload.incomplete_since,
+				"IncompleteSince on Asset Hub should match the RC value"
+			);
+		} else {
+			assert_eq!(
+				pallet_scheduler::IncompleteSince::<T>::get(),
+				ah_incomplete_since,
+				"IncompleteSince on Asset Hub should match the AH value when None from RC"
+			);
+		}
 
 		// Mirror the Agenda conversion in `do_process_scheduler_message` above ^ to construct
 		// expected Agendas. Critically, use the passed agenda call encodings to remove reliance

pallets/ah-migrator/src/society.rs

@@ -214,15 +214,26 @@ pub mod tests {
 				"DefenderVotes map should be empty on the relay chain after migration"
 			);
 
-			assert!(
-				!NextIntakeAt::<T::KusamaConfig>::exists(),
-				"NextIntakeAt should be empty on the relay chain after migration"
-			);
-
-			assert!(
-				!NextChallengeAt::<T::KusamaConfig>::exists(),
-				"NextChallengeAt should be empty on the relay chain after migration"
-			);
+			if let Some(next_challenge_at) = NextChallengeAt::<T::KusamaConfig>::get() {
+				let challenge_period =
+					<T::KusamaConfig as pallet_society::Config>::ChallengePeriod::get();
+				assert_eq!(
+					next_challenge_at, challenge_period,
+					"`next_challenge_at` must be equal to the `ChallengePeriod` if not `None`",
+				);
+			};
+
+			if let Some(next_intake_at) = NextIntakeAt::<T::KusamaConfig>::get() {
+				let rotation_period =
+					<T::KusamaConfig as pallet_society::Config>::VotingPeriod::get()
+						.saturating_add(
+							<T::KusamaConfig as pallet_society::Config>::ClaimPeriod::get(),
+						);
+				assert_eq!(
+					next_intake_at, rotation_period,
+					"`next_intake_at` must be equal to the rotation period if not `None`",
+				);
+			};
 		}
 
 		fn post_check(rc_payload: Self::RcPrePayload, _: Self::AhPrePayload) {

pallets/rc-migrator/src/accounts.rs

@@ -1279,8 +1279,7 @@ pub mod tests {
 			}
 
 			let total_issuance = <T as Config>::Currency::total_issuance();
-			let tracker = RcMigratedBalance::<T>::get();
-			// verify total issuance hasn't changed for any other reason than the migrated funds
+			let tracker = RcMigratedBalanceArchive::<T>::get();
 			assert_eq!(
 				total_issuance,
 				rc_total_issuance_before.saturating_sub(tracker.migrated),

pallets/rc-migrator/src/lib.rs

@@ -786,6 +786,16 @@ pub mod pallet {
 	pub type RcMigratedBalance<T: Config> =
 		StorageValue<_, MigratedBalances<T::Balance>, ValueQuery>;
 
+	/// Helper storage item to store the total balance that should be kept on Relay Chain after
+	/// it is consumed from the `RcMigratedBalance` storage item and sent to the Asset Hub.
+	///
+	/// This let us to take the value from the `RcMigratedBalance` storage item and keep the
+	/// `SignalMigrationFinish` stage to be idempotent while preserving these values for tests and
+	/// later discoveries.
+	#[pallet::storage]
+	pub type RcMigratedBalanceArchive<T: Config> =
+		StorageValue<_, MigratedBalances<T::Balance>, ValueQuery>;
+
 	/// The pending XCM messages.
 	///
 	/// Contains data messages that have been sent to the Asset Hub but not yet confirmed.
@@ -2243,20 +2253,16 @@ pub mod pallet {
 						// 1 read and 1 write for `staking::on_migration_end`;
 						// 1 read and 1 write for `RcMigratedBalance` storage item;
 						// plus one xcm send;
-						T::DbWeight::get().reads_writes(1, 1)
+						T::DbWeight::get().reads_writes(1, 2)
 							.saturating_add(T::RcWeightInfo::send_chunked_xcm_and_track())
 					);
 
 					pallet_staking_async_ah_client::Pallet::<T>::on_migration_end();
 
 					// Send finish message to AH.
 					let data = if RcMigratedBalance::<T>::exists() {
-						let tracker = if cfg!(feature = "std") {
-							// we should keep this value for the tests.
-							RcMigratedBalance::<T>::get()
-						} else {
-							RcMigratedBalance::<T>::take()
-						};
+						let tracker = RcMigratedBalance::<T>::take();
+						RcMigratedBalanceArchive::<T>::put(&tracker);
 						Self::deposit_event(Event::MigratedBalanceConsumed {
 							kept: tracker.kept,
 							migrated: tracker.migrated,

pallets/rc-migrator/src/society.rs

@@ -133,6 +133,31 @@ impl SocietyValues {
 	{
 		use pallet_society::*;
 
+		let next_intake_at = if let Some(next_intake_at) = NextIntakeAt::<T>::take() {
+			let rotation_period = T::VotingPeriod::get().saturating_add(T::ClaimPeriod::get());
+			if next_intake_at != rotation_period {
+				Some(next_intake_at)
+			} else {
+				// current `next_intake_at` is the result of the `on_initialize` execution with
+				// disabled rotation. this may happen if this part of migration is executed twice.
+				None
+			}
+		} else {
+			None
+		};
+		let next_challenge_at = if let Some(next_challenge_at) = NextChallengeAt::<T>::take() {
+			let challenge_period = T::ChallengePeriod::get();
+			if next_challenge_at != challenge_period {
+				Some(next_challenge_at)
+			} else {
+				// current `next_challenge_at` is the result of the `on_initialize` execution with
+				// disabled rotation. this may happen if this part of migration is executed twice.
+				None
+			}
+		} else {
+			None
+		};
+
 		SocietyValues {
 			parameters: Parameters::<T>::take().map(|p| p.into_portable()),
 			pot: Pot::<T>::exists().then(Pot::<T>::take),
@@ -150,8 +175,8 @@ impl SocietyValues {
 				.then(ChallengeRoundCount::<T>::take),
 			defending: Defending::<T>::take()
 				.map(|(a, b, portable_tally)| (a, b, portable_tally.into_portable())),
-			next_intake_at: NextIntakeAt::<T>::take(),
-			next_challenge_at: NextChallengeAt::<T>::take(),
+			next_intake_at,
+			next_challenge_at,
 		}
 	}
 
@@ -824,8 +849,34 @@ pub mod tests {
 					.collect();
 			let defender_votes: Vec<(u32, AccountId32, pallet_society::Vote)> =
 				DefenderVotes::<T::KusamaConfig>::iter().collect();
-			let next_intake_at = NextIntakeAt::<T::KusamaConfig>::get();
-			let next_challenge_at = NextChallengeAt::<T::KusamaConfig>::get();
+
+			let next_intake_at =
+				if let Some(next_intake_at) = NextIntakeAt::<T::KusamaConfig>::get() {
+					let rotation_period =
+						<T::KusamaConfig as pallet_society::Config>::VotingPeriod::get()
+							.saturating_add(
+								<T::KusamaConfig as pallet_society::Config>::ClaimPeriod::get(),
+							);
+					if next_intake_at != rotation_period {
+						Some(next_intake_at)
+					} else {
+						None
+					}
+				} else {
+					None
+				};
+			let next_challenge_at =
+				if let Some(next_challenge_at) = NextChallengeAt::<T::KusamaConfig>::get() {
+					let challenge_period =
+						<T::KusamaConfig as pallet_society::Config>::ChallengePeriod::get();
+					if next_challenge_at != challenge_period {
+						Some(next_challenge_at)
+					} else {
+						None
+					}
+				} else {
+					None
+				};
 
 			RcPrePayload {
 				parameters,
@@ -956,15 +1007,26 @@ pub mod tests {
 				"DefenderVotes map should be empty on the relay chain after migration"
 			);
 
-			assert!(
-				!NextIntakeAt::<T::KusamaConfig>::exists(),
-				"NextIntakeAt should be None on the relay chain after migration"
-			);
-
-			assert!(
-				!NextChallengeAt::<T::KusamaConfig>::exists(),
-				"NextChallengeAt should be None on the relay chain after migration"
-			);
+			if let Some(next_challenge_at) = NextChallengeAt::<T::KusamaConfig>::get() {
+				let challenge_period =
+					<T::KusamaConfig as pallet_society::Config>::ChallengePeriod::get();
+				assert_eq!(
+					next_challenge_at, challenge_period,
+					"`next_challenge_at` must be equal to the `ChallengePeriod` if not `None`",
+				);
+			};
+
+			if let Some(next_intake_at) = NextIntakeAt::<T::KusamaConfig>::get() {
+				let rotation_period =
+					<T::KusamaConfig as pallet_society::Config>::VotingPeriod::get()
+						.saturating_add(
+							<T::KusamaConfig as pallet_society::Config>::ClaimPeriod::get(),
+						);
+				assert_eq!(
+					next_intake_at, rotation_period,
+					"`next_intake_at` must be equal to the rotation period if not `None`",
+				);
+			};
 		}
 	}
 }

pallets/rc-migrator/src/types.rs

@@ -391,6 +391,28 @@ impl<Status: MigrationStatus, Left: TypedGet, Right: Get<Left::Type>> Get<Left::
 	}
 }
 
+/// A value that is `Left::get()` if the migration is finished, otherwise it is `Right::get()`.
+pub struct LeftIfFinished<Status, Left, Right>(PhantomData<(Status, Left, Right)>);
+impl<Status: MigrationStatus, Left: TypedGet, Right: Get<Left::Type>> Get<Left::Type>
+	for LeftIfFinished<Status, Left, Right>
+{
+	fn get() -> Left::Type {
+		Status::is_ongoing().then(|| Left::get()).unwrap_or_else(|| Right::get())
+	}
+}
+
+/// A value that is `Left::get()` if the migration is finished, otherwise it is `Right::get()`.
+pub struct LeftIfPending<Status, Left, Right>(PhantomData<(Status, Left, Right)>);
+impl<Status: MigrationStatus, Left: TypedGet, Right: Get<Left::Type>> Get<Left::Type>
+	for LeftIfPending<Status, Left, Right>
+{
+	fn get() -> Left::Type {
+		(!Status::is_ongoing() && !Status::is_finished())
+			.then(|| Left::get())
+			.unwrap_or_else(|| Right::get())
+	}
+}
+
 /// A weight that is `Weight::MAX` if the migration is ongoing, otherwise it is the [`Inner`]
 /// weight of the [`pallet_fast_unstake::weights::WeightInfo`] trait.
 pub struct MaxOnIdleOrInner<Status, Inner>(PhantomData<(Status, Inner)>);

relay/kusama/src/lib.rs

@@ -1212,11 +1212,23 @@ impl pallet_society::Config for Runtime {
 	type Randomness = pallet_babe::RandomnessFromOneEpochAgo<Runtime>;
 	type GraceStrikes = ConstU32<10>;
 	type PeriodSpend = ConstU128<{ 500 * QUID }>;
-	type VotingPeriod = ConstU32<{ 5 * DAYS }>;
+	type VotingPeriod = pallet_rc_migrator::types::LeftIfPending<
+		RcMigrator,
+		ConstU32<{ 5 * DAYS }>,
+		// disable rotation `on_initialize` during and after migration
+		// { - 10 * DAYS } to avoid the overflow (`VotingPeriod` is summed with `ClaimPeriod`)
+		ConstU32<{ u32::MAX - 10 * DAYS }>,
+	>;
 	type ClaimPeriod = ConstU32<{ 2 * DAYS }>;
 	type MaxLockDuration = ConstU32<{ 36 * 30 * DAYS }>;
 	type FounderSetOrigin = EnsureRoot<AccountId>;
-	type ChallengePeriod = ConstU32<{ 7 * DAYS }>;
+	type ChallengePeriod = pallet_rc_migrator::types::LeftIfPending<
+		RcMigrator,
+		ConstU32<{ 7 * DAYS }>,
+		// disable challenge rotation `on_initialize` during and after migration
+		// { - 10 * DAYS } to make sure we don't overflow
+		ConstU32<{ u32::MAX - 10 * DAYS }>,
+	>;
 	type MaxPayouts = ConstU32<8>;
 	type MaxBids = ConstU32<512>;
 	type PalletId = SocietyPalletId;

system-parachains/asset-hubs/asset-hub-kusama/src/lib.rs

@@ -1536,20 +1536,22 @@ impl pallet_society::Config for Runtime {
 	>;
 	type GraceStrikes = ConstU32<10>;
 	type PeriodSpend = ConstU128<{ 500 * QUID }>;
-	type VotingPeriod = pallet_ah_migrator::LeftOrRight<
+	type VotingPeriod = pallet_ah_migrator::LeftIfFinished<
 		AhMigrator,
-		// disable rotation `on_initialize` during migration
+		ConstU32<{ 5 * RC_DAYS }>,
+		// disable rotation `on_initialize` before and during migration
 		// { - 10 * RC_DAYS } to avoid the overflow (`VotingPeriod` is summed with `ClaimPeriod`)
 		ConstU32<{ u32::MAX - 10 * RC_DAYS }>,
-		ConstU32<{ 5 * RC_DAYS }>,
 	>;
 	type ClaimPeriod = ConstU32<{ 2 * RC_DAYS }>;
 	type MaxLockDuration = ConstU32<{ 36 * 30 * RC_DAYS }>;
 	type FounderSetOrigin = EnsureRoot<AccountId>;
-	type ChallengePeriod = pallet_ah_migrator::LeftOrRight<
+	type ChallengePeriod = pallet_ah_migrator::LeftIfFinished<
 		AhMigrator,
-		ConstU32<{ u32::MAX }>, // disable challenge rotation `on_initialize` during migration
 		ConstU32<{ 7 * RC_DAYS }>,
+		// disable challenge rotation `on_initialize` before and during migration
+		// { - 10 * RC_DAYS } to make sure we don't overflow
+		ConstU32<{ u32::MAX - 10 * RC_DAYS }>,
 	>;
 	type MaxPayouts = ConstU32<8>;
 	type MaxBids = ConstU32<512>;