[AHM] Cache test snapshots and include more tests into CI job

.github/workflows/test.yml

@@ -187,12 +187,17 @@ jobs:
           lz4 -d ah-polkadot.snap.lz4 ah-polkadot.snap
 
       - name: Test Asset Hub Migration (Polkadot)
-        run: cargo test --locked -q -p polkadot-integration-tests-ahm -r pallet_migration_works -- --nocapture
+        run: cargo test --locked -q -p polkadot-integration-tests-ahm -r -- --nocapture
         env:
           RUST_BACKTRACE: 1
           SNAP_RC: "../../polkadot.snap"
           SNAP_AH: "../../ah-polkadot.snap"
           RUST_LOG: "warn"
+      - name: Test Benchmarks Asset Hub Migration (Polkadot)
+        run: cargo test --locked -q -p polkadot-integration-tests-ahm -r --features runtime-benchmarks -- bench_rc bench_ops bench_ah --nocapture
+        env:
+          RUST_BACKTRACE: 1
+          RUST_LOG: "warn"
 
   # TODO: remove if not needed. or reuse for Paseo
   ahm-test-westend:

integration-tests/ahm/src/mock.rs

@@ -16,9 +16,7 @@
 
 use crate::porting_prelude::*;
 
-use asset_hub_polkadot_runtime::{
-	AhMigrator, Block as AssetHubBlock, Runtime as AssetHub, RuntimeEvent as AhRuntimeEvent,
-};
+use asset_hub_polkadot_runtime::{AhMigrator, Runtime as AssetHub, RuntimeEvent as AhRuntimeEvent};
 use codec::Decode;
 use cumulus_primitives_core::{
 	AggregateMessageOrigin as ParachainMessageOrigin, InboundDownwardMessage, ParaId,
@@ -33,46 +31,88 @@ use polkadot_primitives::UpwardMessage;
 use polkadot_runtime::{
 	Block as PolkadotBlock, RcMigrator, Runtime as Polkadot, RuntimeEvent as RcRuntimeEvent,
 };
-use remote_externalities::{Builder, Mode, OfflineConfig, RemoteExternalities};
+use remote_externalities::{Builder, Mode, OfflineConfig};
 use runtime_parachains::{
 	dmp::DownwardMessageQueues,
 	inclusion::{AggregateMessageOrigin as RcMessageOrigin, UmpQueueId},
 };
+use sp_io::TestExternalities;
 use sp_runtime::{BoundedVec, Perbill};
 use std::str::FromStr;
+use tokio::sync::OnceCell;
 use xcm::prelude::*;
 //use frame_support::traits::QueueFootprintQuery; // Only on westend
 
 pub const AH_PARA_ID: ParaId = ParaId::new(1000);
 const LOG_RC: &str = "runtime::relay";
 const LOG_AH: &str = "runtime::asset-hub";
 
+pub enum Chain {
+	Relay,
+	AssetHub,
+}
+
+impl ToString for Chain {
+	fn to_string(&self) -> String {
+		match self {
+			Chain::Relay => "SNAP_RC".to_string(),
+			Chain::AssetHub => "SNAP_AH".to_string(),
+		}
+	}
+}
+
+pub type Snapshot = (Vec<(Vec<u8>, (Vec<u8>, i32))>, sp_core::H256);
+
+static RC_CACHE: OnceCell<Snapshot> = OnceCell::const_new();
+static AH_CACHE: OnceCell<Snapshot> = OnceCell::const_new();
+
 /// Load Relay and AH externalities in parallel.
-pub async fn load_externalities(
-) -> Option<(RemoteExternalities<PolkadotBlock>, RemoteExternalities<AssetHubBlock>)> {
+pub async fn load_externalities() -> Option<(TestExternalities, TestExternalities)> {
 	let (rc, ah) = tokio::try_join!(
-		tokio::spawn(async { remote_ext_test_setup::<PolkadotBlock>("SNAP_RC").await }),
-		tokio::spawn(async { remote_ext_test_setup::<AssetHubBlock>("SNAP_AH").await })
+		tokio::spawn(async { remote_ext_test_setup(Chain::Relay).await }),
+		tokio::spawn(async { remote_ext_test_setup(Chain::AssetHub).await })
 	)
 	.ok()?;
 	Some((rc?, ah?))
 }
 
-pub async fn remote_ext_test_setup<Block: sp_runtime::traits::Block>(
-	env: &str,
-) -> Option<RemoteExternalities<Block>> {
+pub async fn remote_ext_test_setup(chain: Chain) -> Option<TestExternalities> {
 	sp_tracing::try_init_simple();
-	let snap = std::env::var(env).ok()?;
-	let abs = std::path::absolute(snap.clone());
-
-	let ext = Builder::<Block>::default()
-		.mode(Mode::Offline(OfflineConfig { state_snapshot: snap.clone().into() }))
-		.build()
-		.await
-		.map_err(|e| {
-			eprintln!("Could not load from snapshot: {:?}: {:?}", abs, e);
+	log::info!("Checking {} snapshot cache", chain.to_string());
+
+	let cache = match chain {
+		Chain::Relay => &RC_CACHE,
+		Chain::AssetHub => &AH_CACHE,
+	};
+
+	let snapshot = cache
+		.get_or_init(|| async {
+			log::info!("Loading {} snapshot", chain.to_string());
+
+			// Load snapshot.
+			let snap = std::env::var(chain.to_string()).ok().expect("Env var not set");
+			let abs = std::path::absolute(snap.clone());
+
+			let ext = Builder::<PolkadotBlock>::default()
+				.mode(Mode::Offline(OfflineConfig { state_snapshot: snap.clone().into() }))
+				.build()
+				.await
+				.map_err(|e| {
+					eprintln!("Could not load from snapshot: {:?}: {:?}", abs, e);
+				})
+				.unwrap();
+
+			// `RemoteExternalities` and `TestExternalities` types cannot be cloned so we need to
+			// convert them to raw snapshot and store it in the cache.
+			ext.inner_ext.into_raw_snapshot()
 		})
-		.unwrap();
+		.await;
+
+	let ext = TestExternalities::from_raw_snapshot(
+		snapshot.0.clone(),
+		snapshot.1.clone(),
+		sp_storage::StateVersion::V1,
+	);
 
 	Some(ext)
 }
@@ -250,7 +290,7 @@ fn sanity_check_xcm<Call: Decode>(msg: &[u8]) {
 // stage. Otherwise, the `AccountsMigrationInit` stage will be set bypassing the `Scheduled` stage.
 // The `Scheduled` stage is tested separately by the `scheduled_migration_works` test.
 pub fn set_initial_migration_stage(
-	relay_chain: &mut RemoteExternalities<PolkadotBlock>,
+	relay_chain: &mut TestExternalities,
 ) -> RcMigrationStageOf<Polkadot> {
 	let stage = relay_chain.execute_with(|| {
 		let stage = if let Ok(stage) = std::env::var("START_STAGE") {
@@ -272,9 +312,7 @@ pub fn set_initial_migration_stage(
 // both the DMP messages sent from the relay chain to asset hub, which will be used to perform the
 // migration, and the relay chain payload, which will be used to check the correctness of the
 // migration process.
-pub fn rc_migrate(
-	relay_chain: &mut RemoteExternalities<PolkadotBlock>,
-) -> Vec<InboundDownwardMessage> {
+pub fn rc_migrate(relay_chain: &mut TestExternalities) -> Vec<InboundDownwardMessage> {
 	// AH parachain ID
 	let para_id = ParaId::from(1000);
 
@@ -317,10 +355,7 @@ pub fn rc_migrate(
 // Processes all the pending DMP messages in the AH message queue to complete the pallet
 // migration. Uses the relay chain pre-migration payload to check the correctness of the
 // migration once completed.
-pub fn ah_migrate(
-	asset_hub: &mut RemoteExternalities<AssetHubBlock>,
-	dmp_messages: Vec<InboundDownwardMessage>,
-) {
+pub fn ah_migrate(asset_hub: &mut TestExternalities, dmp_messages: Vec<InboundDownwardMessage>) {
 	// Inject the DMP messages into the Asset Hub
 	asset_hub.execute_with(|| {
 		let mut fp =

integration-tests/ahm/src/tests.rs

@@ -60,11 +60,11 @@ use pallet_rc_migrator::{
 	RcMigrationStage as RcMigrationStageStorage,
 };
 use polkadot_primitives::UpwardMessage;
-use polkadot_runtime::{Block as PolkadotBlock, RcMigrator, Runtime as Polkadot};
+use polkadot_runtime::{RcMigrator, Runtime as Polkadot};
 use polkadot_runtime_common::{paras_registrar, slots as pallet_slots};
-use remote_externalities::RemoteExternalities;
 use runtime_parachains::dmp::DownwardMessageQueues;
 use sp_core::crypto::Ss58Codec;
+use sp_io::TestExternalities;
 use sp_runtime::{AccountId32, DispatchError, TokenError};
 use std::{
 	collections::{BTreeMap, VecDeque},
@@ -151,6 +151,7 @@ pub type AhPolkadotChecks = (
 #[cfg(not(feature = "ahm-polkadot"))]
 pub type AhPolkadotChecks = ();
 
+#[ignore] // we use the equivalent [migration_works_time] test instead
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn pallet_migration_works() {
 	let (mut rc, mut ah) = load_externalities().await.unwrap();
@@ -194,7 +195,7 @@ async fn pallet_migration_works() {
 	run_check(|| AhChecks::post_check(rc_pre.unwrap(), ah_pre.unwrap()), &mut ah);
 }
 
-fn run_check<R, B: BlockT>(f: impl FnOnce() -> R, ext: &mut RemoteExternalities<B>) -> Option<R> {
+fn run_check<R>(f: impl FnOnce() -> R, ext: &mut TestExternalities) -> Option<R> {
 	if std::env::var("START_STAGE").is_err() {
 		Some(ext.execute_with(|| f()))
 	} else {
@@ -205,7 +206,7 @@ fn run_check<R, B: BlockT>(f: impl FnOnce() -> R, ext: &mut RemoteExternalities<
 #[cfg(not(feature = "ahm-westend"))] // No auctions on Westend
 #[tokio::test]
 async fn num_leases_to_ending_block_works_simple() {
-	let mut rc = remote_ext_test_setup::<PolkadotBlock>("SNAP_RC").await.unwrap();
+	let mut rc = remote_ext_test_setup(Chain::Relay).await.unwrap();
 	let f = |now: BlockNumberFor<Polkadot>, num_leases: u32| {
 		frame_system::Pallet::<Polkadot>::set_block_number(now);
 		pallet_rc_migrator::crowdloan::num_leases_to_ending_block::<Polkadot>(num_leases)
@@ -324,7 +325,7 @@ async fn print_sovereign_account_translation() {
 async fn print_accounts_statistics() {
 	use frame_system::Account as SystemAccount;
 
-	let mut rc = remote_ext_test_setup::<PolkadotBlock>("SNAP_RC").await.unwrap();
+	let mut rc = remote_ext_test_setup(Chain::Relay).await.unwrap();
 
 	let mut total_counts = std::collections::HashMap::new();
 
@@ -387,8 +388,7 @@ fn ah_account_migration_weight() {
 	}
 }
 
-#[ignore] // Slow
-#[tokio::test(flavor = "current_thread")]
+#[tokio::test]
 async fn migration_works_time() {
 	let Some((mut rc, mut ah)) = load_externalities().await else { return };
 
@@ -488,7 +488,7 @@ async fn migration_works_time() {
 	);
 }
 
-#[tokio::test(flavor = "current_thread")]
+#[tokio::test]
 async fn scheduled_migration_works() {
 	let Some((mut rc, mut ah)) = load_externalities().await else { return };