Make xcm-simulator async with more tests (#3618)

* make xcm simulator async

* fmt

* nightly fmt

* Update xcm/xcm-simulator/example/src/parachain.rs

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>

* accept keith review suggestion

* master.into

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>

xcm/xcm-simulator/example/src/lib.rs

@@ -17,9 +17,12 @@
 mod parachain;
 mod relay_chain;
 
+use polkadot_parachain::primitives::Id as ParaId;
+use sp_runtime::traits::AccountIdConversion;
 use xcm_simulator::{decl_test_network, decl_test_parachain, decl_test_relay_chain};
 
 pub const ALICE: sp_runtime::AccountId32 = sp_runtime::AccountId32::new([0u8; 32]);
+pub const INITIAL_BALANCE: u128 = 1_000_000_000;
 
 decl_test_parachain! {
 	pub struct ParaA {
@@ -57,7 +60,9 @@ decl_test_network! {
 	}
 }
 
-pub const INITIAL_BALANCE: u128 = 1_000_000_000;
+pub fn para_account_id(id: u32) -> relay_chain::AccountId {
+	ParaId::from(id).into_account()
+}
 
 pub fn para_ext(para_id: u32) -> sp_io::TestExternalities {
 	use parachain::{MsgQueue, Runtime, System};
@@ -81,9 +86,11 @@ pub fn relay_ext() -> sp_io::TestExternalities {
 
 	let mut t = frame_system::GenesisConfig::default().build_storage::<Runtime>().unwrap();
 
-	pallet_balances::GenesisConfig::<Runtime> { balances: vec![(ALICE, INITIAL_BALANCE)] }
-		.assimilate_storage(&mut t)
-		.unwrap();
+	pallet_balances::GenesisConfig::<Runtime> {
+		balances: vec![(ALICE, INITIAL_BALANCE), (para_account_id(1), INITIAL_BALANCE)],
+	}
+	.assimilate_storage(&mut t)
+	.unwrap();
 
 	let mut ext = sp_io::TestExternalities::new(t);
 	ext.execute_with(|| System::set_block_number(1));
@@ -98,10 +105,21 @@ mod tests {
 	use super::*;
 
 	use codec::Encode;
-	use frame_support::assert_ok;
+	use frame_support::{assert_ok, weights::Weight};
 	use xcm::latest::prelude::*;
 	use xcm_simulator::TestExt;
 
+	// Helper function for forming buy execution message
+	fn buy_execution<C>(fees: impl Into<MultiAsset>, debt: Weight) -> Order<C> {
+		Order::BuyExecution {
+			fees: fees.into(),
+			weight: 0,
+			debt,
+			halt_on_error: false,
+			instructions: vec![],
+		}
+	}
+
 	#[test]
 	fn dmp() {
 		MockNet::reset();
@@ -187,22 +205,123 @@ mod tests {
 	fn reserve_transfer() {
 		MockNet::reset();
 
+		let withdraw_amount = 123;
+		let max_weight_for_execution = 3;
+
 		Relay::execute_with(|| {
 			assert_ok!(RelayChainPalletXcm::reserve_transfer_assets(
 				relay_chain::Origin::signed(ALICE),
 				Box::new(X1(Parachain(1)).into()),
 				Box::new(X1(AccountId32 { network: Any, id: ALICE.into() }).into()),
-				(Here, 123).into(),
+				(Here, withdraw_amount).into(),
 				0,
-				3,
+				max_weight_for_execution,
 			));
+			assert_eq!(
+				parachain::Balances::free_balance(&para_account_id(1)),
+				INITIAL_BALANCE + withdraw_amount
+			);
 		});
 
 		ParaA::execute_with(|| {
 			// free execution, full amount received
 			assert_eq!(
 				pallet_balances::Pallet::<parachain::Runtime>::free_balance(&ALICE),
-				INITIAL_BALANCE + 123
+				INITIAL_BALANCE + withdraw_amount
+			);
+		});
+	}
+
+	/// Scenario:
+	/// A parachain transfers funds on the relay chain to another parachain account.
+	///
+	/// Asserts that the parachain accounts are updated as expected.
+	#[test]
+	fn withdraw_and_deposit() {
+		MockNet::reset();
+
+		let send_amount = 10;
+		let weight_for_execution = 3 * relay_chain::BaseXcmWeight::get();
+
+		ParaA::execute_with(|| {
+			let message = WithdrawAsset {
+				assets: (Here, send_amount).into(),
+				effects: vec![
+					buy_execution((Here, send_amount), weight_for_execution),
+					Order::DepositAsset {
+						assets: All.into(),
+						max_assets: 1,
+						beneficiary: Parachain(2).into(),
+					},
+				],
+			};
+			// Send withdraw and deposit
+			assert_ok!(ParachainPalletXcm::send_xcm(Here, Parent.into(), message.clone()));
+		});
+
+		Relay::execute_with(|| {
+			assert_eq!(
+				relay_chain::Balances::free_balance(para_account_id(1)),
+				INITIAL_BALANCE - send_amount
+			);
+			assert_eq!(relay_chain::Balances::free_balance(para_account_id(2)), send_amount);
+		});
+	}
+
+	/// Scenario:
+	/// A parachain wants to be notified that a transfer worked correctly.
+	/// It sends a `QueryHolding` after the deposit to get notified on success.
+	///
+	/// Asserts that the balances are updated correctly and the expected XCM is sent.
+	#[test]
+	fn query_holding() {
+		MockNet::reset();
+
+		let send_amount = 10;
+		let weight_for_execution = 3 * relay_chain::BaseXcmWeight::get();
+		let query_id_set = 1234;
+
+		// Send a message which fully succeeds on the relay chain
+		ParaA::execute_with(|| {
+			let message = WithdrawAsset {
+				assets: (Here, send_amount).into(),
+				effects: vec![
+					buy_execution((Here, send_amount), weight_for_execution),
+					Order::DepositAsset {
+						assets: All.into(),
+						max_assets: 1,
+						beneficiary: Parachain(2).into(),
+					},
+					Order::QueryHolding {
+						query_id: query_id_set,
+						dest: Parachain(1).into(),
+						assets: All.into(),
+					},
+				],
+			};
+			// Send withdraw and deposit with query holding
+			assert_ok!(ParachainPalletXcm::send_xcm(Here, Parent.into(), message.clone(),));
+		});
+
+		// Check that transfer was executed
+		Relay::execute_with(|| {
+			// Withdraw executed
+			assert_eq!(
+				relay_chain::Balances::free_balance(para_account_id(1)),
+				INITIAL_BALANCE - send_amount
+			);
+			// Deposit executed
+			assert_eq!(relay_chain::Balances::free_balance(para_account_id(2)), send_amount);
+		});
+
+		// Check that QueryResponse message was received
+		ParaA::execute_with(|| {
+			assert_eq!(
+				parachain::MsgQueue::received_dmp(),
+				vec![QueryResponse {
+					query_id: query_id_set,
+					response: Response::Assets(MultiAssets::new())
+				}]
 			);
 		});
 	}

xcm/xcm-simulator/example/src/parachain.rs

@@ -166,6 +166,11 @@ pub mod mock_msg_queue {
 	#[pallet::getter(fn parachain_id)]
 	pub(super) type ParachainId<T: Config> = StorageValue<_, ParaId, ValueQuery>;
 
+	#[pallet::storage]
+	#[pallet::getter(fn received_dmp)]
+	/// A queue of received DMP messages
+	pub(super) type ReceivedDmp<T: Config> = StorageValue<_, Vec<Xcm<T::Call>>, ValueQuery>;
+
 	impl<T: Config> Get<ParaId> for Pallet<T> {
 		fn get() -> ParaId {
 			Self::parachain_id()
@@ -266,7 +271,8 @@ pub mod mock_msg_queue {
 						Self::deposit_event(Event::UnsupportedVersion(id));
 					},
 					Ok(Ok(x)) => {
-						let outcome = T::XcmExecutor::execute_xcm(Parent.into(), x, limit);
+						let outcome = T::XcmExecutor::execute_xcm(Parent.into(), x.clone(), limit);
+						<ReceivedDmp<T>>::append(x);
 						Self::deposit_event(Event::ExecutedDownward(id, outcome));
 					},
 				}

xcm/xcm-simulator/src/lib.rs

@@ -21,7 +21,7 @@ pub use paste;
 
 pub use frame_support::{traits::Get, weights::Weight};
 pub use sp_io::TestExternalities;
-pub use sp_std::{cell::RefCell, marker::PhantomData};
+pub use sp_std::{cell::RefCell, collections::vec_deque::VecDeque, marker::PhantomData};
 
 pub use polkadot_core_primitives::BlockNumber as RelayBlockNumber;
 pub use polkadot_parachain::primitives::{
@@ -36,9 +36,24 @@ pub use xcm::{latest::prelude::*, VersionedXcm};
 pub use xcm_executor::XcmExecutor;
 
 pub trait TestExt {
+	/// Initialize the test environment.
 	fn new_ext() -> sp_io::TestExternalities;
+	/// Resets the state of the test environment.
 	fn reset_ext();
-	fn execute_with<R>(execute: impl FnOnce() -> R) -> R;
+	/// Execute code in the context of the test externalities, without automatic
+	/// message processing. All messages in the message buses can be processed
+	/// by calling `Self::dispatch_xcm_buses()`.
+	fn execute_without_dispatch<R>(execute: impl FnOnce() -> R) -> R;
+	/// Process all messages in the message buses
+	fn dispatch_xcm_buses();
+	/// Execute some code in the context of the test externalities, with
+	/// automatic message processing.
+	/// Messages are dispatched once the passed closure completes.
+	fn execute_with<R>(execute: impl FnOnce() -> R) -> R {
+		let result = Self::execute_without_dispatch(execute);
+		Self::dispatch_xcm_buses();
+		result
+	}
 }
 
 pub enum MessageKind {
@@ -162,13 +177,31 @@ macro_rules! __impl_ext {
 				$ext_name.with(|v| *v.borrow_mut() = $new_ext);
 			}
 
-			fn execute_with<R>(execute: impl FnOnce() -> R) -> R {
+			fn execute_without_dispatch<R>(execute: impl FnOnce() -> R) -> R {
 				$ext_name.with(|v| v.borrow_mut().execute_with(execute))
 			}
+
+			fn dispatch_xcm_buses() {
+				while exists_messages_in_any_bus() {
+					if let Err(xcm_error) = process_relay_messages() {
+						panic!("Relay chain XCM execution failure: {:?}", xcm_error);
+					}
+					if let Err(xcm_error) = process_para_messages() {
+						panic!("Parachain XCM execution failure: {:?}", xcm_error);
+					}
+				}
+			}
 		}
 	};
 }
 
+thread_local! {
+	pub static PARA_MESSAGE_BUS: RefCell<VecDeque<(ParaId, MultiLocation, Xcm<()>)>>
+		= RefCell::new(VecDeque::new());
+	pub static RELAY_MESSAGE_BUS: RefCell<VecDeque<(MultiLocation, Xcm<()>)>>
+		= RefCell::new(VecDeque::new());
+}
+
 #[macro_export]
 macro_rules! decl_test_network {
 	(
@@ -181,37 +214,101 @@ macro_rules! decl_test_network {
 
 		impl $name {
 			pub fn reset() {
-				use $crate::TestExt;
-
+				use $crate::{TestExt, VecDeque};
+				// Reset relay chain message bus
+				$crate::RELAY_MESSAGE_BUS.with(|b| b.replace(VecDeque::new()));
+				// Reset parachain message bus
+				$crate::PARA_MESSAGE_BUS.with(|b| b.replace(VecDeque::new()));
 				<$relay_chain>::reset_ext();
 				$( <$parachain>::reset_ext(); )*
 			}
 		}
 
-		/// XCM router for parachain.
-		pub struct ParachainXcmRouter<T>($crate::PhantomData<T>);
+		/// Check if any messages exist in either message bus
+		fn exists_messages_in_any_bus() -> bool {
+			use $crate::{RELAY_MESSAGE_BUS, PARA_MESSAGE_BUS};
+			let no_relay_messages_left = RELAY_MESSAGE_BUS.with(|b| b.borrow().is_empty());
+			let no_parachain_messages_left = PARA_MESSAGE_BUS.with(|b| b.borrow().is_empty());
+			!(no_relay_messages_left && no_parachain_messages_left)
+		}
 
-		impl<T: $crate::Get<$crate::ParaId>> $crate::SendXcm for ParachainXcmRouter<T> {
-			fn send_xcm(destination: $crate::MultiLocation, message: $crate::Xcm<()>) -> $crate::XcmResult {
-				use $crate::{UmpSink, XcmpMessageHandlerT};
+		/// Process all messages originating from parachains.
+		fn process_para_messages() -> $crate::XcmResult {
+			use $crate::{UmpSink, XcmpMessageHandlerT};
 
+			while let Some((para_id, destination, message)) = $crate::PARA_MESSAGE_BUS.with(
+				|b| b.borrow_mut().pop_front()) {
 				match destination.interior() {
 					$crate::Junctions::Here if destination.parent_count() == 1 => {
 						let encoded = $crate::encode_xcm(message, $crate::MessageKind::Ump);
-						let _ = <$relay_chain>::process_upward_message(
-							T::get(), &encoded[..],
+						let r = <$relay_chain>::process_upward_message(
+							para_id, &encoded[..],
 							$crate::Weight::max_value(),
 						);
-						Ok(())
+						if let Err((id, required)) = r {
+							return Err($crate::XcmError::WeightLimitReached(required));
+						}
 					},
 					$(
 						$crate::X1($crate::Parachain(id)) if *id == $para_id && destination.parent_count() == 1 => {
 							let encoded = $crate::encode_xcm(message, $crate::MessageKind::Xcmp);
-							let messages = vec![(T::get(), 1, &encoded[..])];
-							let _ = <$parachain>::handle_xcmp_messages(
+							let messages = vec![(para_id, 1, &encoded[..])];
+							let _weight = <$parachain>::handle_xcmp_messages(
 								messages.into_iter(),
 								$crate::Weight::max_value(),
 							);
+						},
+					)*
+					_ => {
+						return Err($crate::XcmError::CannotReachDestination(destination, message));
+					}
+				}
+			}
+
+			Ok(())
+		}
+
+		/// Process all messages originating from the relay chain.
+		fn process_relay_messages() -> $crate::XcmResult {
+			use $crate::DmpMessageHandlerT;
+
+			while let Some((destination, message)) = $crate::RELAY_MESSAGE_BUS.with(
+				|b| b.borrow_mut().pop_front()) {
+				match destination.interior() {
+					$(
+						$crate::X1($crate::Parachain(id)) if *id == $para_id && destination.parent_count() == 0 => {
+							let encoded = $crate::encode_xcm(message, $crate::MessageKind::Dmp);
+							// NOTE: RelayChainBlockNumber is hard-coded to 1
+							let messages = vec![(1, encoded)];
+							let _weight = <$parachain>::handle_dmp_messages(
+								messages.into_iter(), $crate::Weight::max_value(),
+							);
+						},
+					)*
+					_ => return Err($crate::XcmError::SendFailed("Only sends to children parachain.")),
+				}
+			}
+
+			Ok(())
+		}
+
+		/// XCM router for parachain.
+		pub struct ParachainXcmRouter<T>($crate::PhantomData<T>);
+
+		impl<T: $crate::Get<$crate::ParaId>> $crate::SendXcm for ParachainXcmRouter<T> {
+			fn send_xcm(destination: $crate::MultiLocation, message: $crate::Xcm<()>) -> $crate::XcmResult {
+				use $crate::{UmpSink, XcmpMessageHandlerT};
+
+				match destination.interior() {
+					$crate::Junctions::Here if destination.parent_count() == 1 => {
+						$crate::PARA_MESSAGE_BUS.with(
+							|b| b.borrow_mut().push_back((T::get(), destination, message)));
+						Ok(())
+					},
+					$(
+						$crate::X1($crate::Parachain(id)) if *id == $para_id && destination.parent_count() == 1 => {
+							$crate::PARA_MESSAGE_BUS.with(
+								|b| b.borrow_mut().push_back((T::get(), destination, message)));
 							Ok(())
 						},
 					)*
@@ -229,11 +326,8 @@ macro_rules! decl_test_network {
 				match destination.interior() {
 					$(
 						$crate::X1($crate::Parachain(id)) if *id == $para_id && destination.parent_count() == 0 => {
-							let encoded = $crate::encode_xcm(message, $crate::MessageKind::Dmp);
-							let messages = vec![(1, encoded)];
-							let _ = <$parachain>::handle_dmp_messages(
-								messages.into_iter(), $crate::Weight::max_value(),
-							);
+							$crate::RELAY_MESSAGE_BUS.with(
+								|b| b.borrow_mut().push_back((destination, message)));
 							Ok(())
 						},
 					)*