From 27811dd0cba4eeb5113e9072e3cc816dd5ee7283 Mon Sep 17 00:00:00 2001
From: Branislav Kontur <bkontur@gmail.com>
Date: Fri, 7 Jun 2024 16:09:31 +0200
Subject: [PATCH] Fix `reap_identity` benchmark

---
 relay/polkadot/src/impls.rs | 30 ++++++++++++++++++++++++++++++
 relay/polkadot/src/lib.rs   |  4 ++++
 2 files changed, 34 insertions(+)

diff --git a/relay/polkadot/src/impls.rs b/relay/polkadot/src/impls.rs
index 0f72c0cec9..5f3c39f328 100644
--- a/relay/polkadot/src/impls.rs
+++ b/relay/polkadot/src/impls.rs
@@ -216,3 +216,33 @@ where
 		I::ensure_successful(asset_kind)
 	}
 }
+
+#[cfg(feature = "runtime-benchmarks")]
+pub mod benchmarks {
+	use super::{xcm_config::CheckAccount, ExistentialDeposit};
+	use crate::Balances;
+	use frame_support::{
+		dispatch::RawOrigin,
+		traits::{Currency, EnsureOrigin},
+	};
+
+	pub struct InitializeReaperForBenchmarking<A, E>(core::marker::PhantomData<(A, E)>);
+	impl<A, O: Into<Result<RawOrigin<A>, O>> + From<RawOrigin<A>>, E: EnsureOrigin<O>>
+		EnsureOrigin<O> for InitializeReaperForBenchmarking<A, E>
+	{
+		type Success = E::Success;
+
+		fn try_origin(o: O) -> Result<E::Success, O> {
+			E::try_origin(o)
+		}
+
+		#[cfg(feature = "runtime-benchmarks")]
+		fn try_successful_origin() -> Result<O, ()> {
+			// initialize the XCM Check Account with the existential deposit
+			Balances::make_free_balance_be(&CheckAccount::get(), ExistentialDeposit::get());
+
+			// call the real implementation
+			E::try_successful_origin()
+		}
+	}
+}
diff --git a/relay/polkadot/src/lib.rs b/relay/polkadot/src/lib.rs
index 66496b7336..d0b8ed41a4 100644
--- a/relay/polkadot/src/lib.rs
+++ b/relay/polkadot/src/lib.rs
@@ -860,7 +860,11 @@ impl pallet_identity::Config for Runtime {
 impl identity_migrator::Config for Runtime {
 	type RuntimeEvent = RuntimeEvent;
 	// To be updated to `EnsureSigned` once the parachain is producing blocks.
+	#[cfg(not(feature = "runtime-benchmarks"))]
 	type Reaper = EnsureRoot<AccountId>;
+	#[cfg(feature = "runtime-benchmarks")]
+	type Reaper =
+		impls::benchmarks::InitializeReaperForBenchmarking<AccountId, EnsureRoot<AccountId>>;
 	type ReapIdentityHandler = ToParachainIdentityReaper<Runtime, Self::AccountId>;
 	type WeightInfo = weights::runtime_common_identity_migrator::WeightInfo<Runtime>;
 }