[xcm-emulator] Add on_initialize and on_finalize hooks

cumulus/parachains/integration-tests/emulated/common/src/lib.rs

@@ -32,11 +32,11 @@ pub use pallet_balances;
 
 // Cumulus
 pub use cumulus_pallet_xcmp_queue;
-pub use xcm_emulator::Chain;
 use xcm_emulator::{
 	decl_test_bridges, decl_test_networks, decl_test_parachains, decl_test_relay_chains,
 	decl_test_sender_receiver_accounts_parameter_types, DefaultMessageProcessor,
 };
+pub use xcm_emulator::{Chain, OnHooks};
 
 // Polkadot
 pub use pallet_xcm;
@@ -58,7 +58,8 @@ decl_test_relay_chains! {
 			Balances: westend_runtime::Balances,
 			Treasury: westend_runtime::Treasury,
 			AssetRate: westend_runtime::AssetRate,
-		}
+		},
+		hooks = (),
 	},
 	#[api_version(8)]
 	pub struct Rococo {
@@ -74,7 +75,8 @@ decl_test_relay_chains! {
 			Sudo: rococo_runtime::Sudo,
 			Balances: rococo_runtime::Balances,
 			Hrmp: rococo_runtime::Hrmp,
-		}
+		},
+		hooks = (),
 	},
 	#[api_version(8)]
 	pub struct Wococo {
@@ -89,7 +91,8 @@ decl_test_relay_chains! {
 			XcmPallet: rococo_runtime::XcmPallet,
 			Sudo: rococo_runtime::Sudo,
 			Balances: rococo_runtime::Balances,
-		}
+		},
+		hooks = (),
 	}
 }
 
@@ -114,7 +117,9 @@ decl_test_parachains! {
 			ForeignAssets: asset_hub_westend_runtime::ForeignAssets,
 			PoolAssets: asset_hub_westend_runtime::PoolAssets,
 			AssetConversion: asset_hub_westend_runtime::AssetConversion,
-		}
+		},
+		hooks = (),
+
 	},
 	pub struct PenpalWestendA {
 		genesis = penpal::genesis(penpal::PARA_ID_A),
@@ -132,7 +137,8 @@ decl_test_parachains! {
 			PolkadotXcm: penpal_runtime::PolkadotXcm,
 			Assets: penpal_runtime::Assets,
 			Balances: penpal_runtime::Balances,
-		}
+		},
+		hooks = (),
 	},
 	// Rococo Parachains
 	pub struct BridgeHubRococo {
@@ -150,7 +156,8 @@ decl_test_parachains! {
 		pallets = {
 			PolkadotXcm: bridge_hub_rococo_runtime::PolkadotXcm,
 			Balances: bridge_hub_rococo_runtime::Balances,
-		}
+		},
+		hooks = (),
 	},
 	// AssetHubRococo
 	pub struct AssetHubRococo {
@@ -172,7 +179,8 @@ decl_test_parachains! {
 			PoolAssets: asset_hub_rococo_runtime::PoolAssets,
 			AssetConversion: asset_hub_rococo_runtime::AssetConversion,
 			Balances: asset_hub_rococo_runtime::Balances,
-		}
+		},
+		hooks = (),
 	},
 	pub struct PenpalRococoA {
 		genesis = penpal::genesis(penpal::PARA_ID_A),
@@ -189,7 +197,8 @@ decl_test_parachains! {
 		pallets = {
 			PolkadotXcm: penpal_runtime::PolkadotXcm,
 			Assets: penpal_runtime::Assets,
-		}
+		},
+		hooks = (),
 	},
 	pub struct PenpalRococoB {
 		genesis = penpal::genesis(penpal::PARA_ID_B),
@@ -206,7 +215,8 @@ decl_test_parachains! {
 		pallets = {
 			PolkadotXcm: penpal_runtime::PolkadotXcm,
 			Assets: penpal_runtime::Assets,
-		}
+		},
+		hooks = (),
 	},
 	// Wococo Parachains
 	pub struct BridgeHubWococo {
@@ -224,7 +234,8 @@ decl_test_parachains! {
 		},
 		pallets = {
 			PolkadotXcm: bridge_hub_rococo_runtime::PolkadotXcm,
-		}
+		},
+		hooks = (),
 	},
 	pub struct AssetHubWococo {
 		genesis = asset_hub_wococo::genesis(),
@@ -246,7 +257,8 @@ decl_test_parachains! {
 			PoolAssets: asset_hub_rococo_runtime::PoolAssets,
 			AssetConversion: asset_hub_rococo_runtime::AssetConversion,
 			Balances: asset_hub_rococo_runtime::Balances,
-		}
+		},
+		hooks = (),
 	}
 }
 

cumulus/xcm/xcm-emulator/src/lib.rs

@@ -60,7 +60,6 @@ pub use xcm::v3::prelude::{
 	XcmHash, X1,
 };
 pub use xcm_executor::traits::ConvertLocation;
-
 pub type AccountIdOf<T> = <T as frame_system::Config>::AccountId;
 
 thread_local! {
@@ -210,12 +209,25 @@ pub trait NetworkComponent {
 	}
 }
 
+pub trait OnHooks {
+	fn on_initialize(n: BlockNumber) -> Weight;
+	fn on_finalize(n: BlockNumber);
+}
+
+impl OnHooks for () {
+	fn on_initialize(_n: BlockNumber) -> Weight {
+		Weight::default()
+	}
+	fn on_finalize(_n: BlockNumber) {}
+}
+
 pub trait Chain: TestExt + NetworkComponent {
 	type Runtime: SystemConfig;
 	type RuntimeCall;
 	type RuntimeOrigin;
 	type RuntimeEvent;
 	type System;
+	type Hooks: OnHooks;
 
 	fn account_id_of(seed: &str) -> AccountId {
 		helpers::get_account_id_from_seed::<sr25519::Public>(seed)
@@ -350,7 +362,8 @@ macro_rules! decl_test_relay_chains {
 				},
 				pallets = {
 					$($pallet_name:ident: $pallet_path:path,)*
-				}
+				},
+				hooks = $on_hooks:ident,
 			}
 		),
 		+
@@ -360,12 +373,15 @@ macro_rules! decl_test_relay_chains {
 			#[derive(Clone)]
 			pub struct $name;
 
+			use $crate::NetworkComponent;
+
 			impl $crate::Chain for $name {
 				type Runtime = $runtime::Runtime;
 				type RuntimeCall = $runtime::RuntimeCall;
 				type RuntimeOrigin = $runtime::RuntimeOrigin;
 				type RuntimeEvent = $runtime::RuntimeEvent;
 				type System = $crate::SystemPallet::<Self::Runtime>;
+				type Hooks = $on_hooks;
 
 				fn account_data_of(account: $crate::AccountIdOf<Self::Runtime>) -> $crate::AccountData<$crate::Balance> {
 					<Self as $crate::TestExt>::ext_wrapper(|| $crate::SystemPallet::<Self::Runtime>::account(account).data.into())
@@ -575,7 +591,8 @@ macro_rules! decl_test_parachains {
 				},
 				pallets = {
 					$($pallet_name:ident: $pallet_path:path,)*
-				}
+				},
+				hooks = $on_hooks:ident,
 			}
 		),
 		+
@@ -591,6 +608,7 @@ macro_rules! decl_test_parachains {
 				type RuntimeOrigin = $runtime::RuntimeOrigin;
 				type RuntimeEvent = $runtime::RuntimeEvent;
 				type System = $crate::SystemPallet::<Self::Runtime>;
+				type Hooks = $on_hooks;
 
 				fn account_data_of(account: $crate::AccountIdOf<Self::Runtime>) -> $crate::AccountData<$crate::Balance> {
 					<Self as $crate::TestExt>::ext_wrapper(|| $crate::SystemPallet::<Self::Runtime>::account(account).data.into())
@@ -645,7 +663,6 @@ macro_rules! decl_test_parachains {
 							.clone()
 						);
 						<Self as Chain>::System::initialize(&block_number, &parent_head_data.hash(), &Default::default());
-						<<Self as Parachain>::ParachainSystem as Hooks<$crate::BlockNumber>>::on_initialize(block_number);
 
 						let _ = <Self as Parachain>::ParachainSystem::set_validation_data(
 							<Self as Chain>::RuntimeOrigin::none(),
@@ -665,7 +682,6 @@ macro_rules! decl_test_parachains {
 					Self::set_last_head();
 				}
 
-
 				fn set_last_head() {
 					use $crate::{Chain, Encode, HeadData, Network, NetworkComponent, Parachain, TestExt};