From a010dc878707db1a4670266ce9823cb6783ca76b Mon Sep 17 00:00:00 2001
From: Aleksandar Brayanov <aleksandar@moonsonglab.com>
Date: Fri, 20 Dec 2024 12:58:32 +0000
Subject: [PATCH] feat: added benchmark helper for pallet_data_preservers

---
 pallets/data-preservers/src/benchmarks.rs | 15 +++++++++++++++
 pallets/data-preservers/src/lib.rs        |  7 +++++++
 runtime/dancebox/src/lib.rs               |  2 ++
 runtime/flashbox/src/lib.rs               |  3 +++
 solo-chains/runtime/dancelight/src/lib.rs |  6 +++---
 5 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/pallets/data-preservers/src/benchmarks.rs b/pallets/data-preservers/src/benchmarks.rs
index 95560050f..02bb2ae75 100644
--- a/pallets/data-preservers/src/benchmarks.rs
+++ b/pallets/data-preservers/src/benchmarks.rs
@@ -36,6 +36,17 @@ use {
     tp_traits::{ParaId, StorageDeposit},
 };
 
+/// Trait describing factory function for para_id.
+pub trait ArgumentFactory<ParaId> {
+    /// Factory function reserving parachain
+    fn reserve_para_id(para_id: ParaId) -> ();
+}
+
+/// Dummy implementation for when the factory is not needed.
+impl<ParaId> ArgumentFactory<ParaId> for () {
+    fn reserve_para_id(_para_id: ParaId) -> () {}
+}
+
 macro_rules! bset {
     ( $($value:expr),* $(,)? ) => {
         {
@@ -333,6 +344,8 @@ mod benchmarks {
         Pallet::<T>::create_profile(RawOrigin::Signed(caller.clone()).into(), profile)
             .expect("to create profile");
 
+        T::BenchmarkHelper::reserve_para_id(para_id);
+
         let origin = T::AssignmentOrigin::try_successful_origin(&para_id).unwrap();
 
         #[extrinsic_call]
@@ -366,6 +379,8 @@ mod benchmarks {
         Pallet::<T>::create_profile(RawOrigin::Signed(caller.clone()).into(), profile)
             .expect("to create profile");
 
+        T::BenchmarkHelper::reserve_para_id(para_id);
+
         let origin = T::AssignmentOrigin::try_successful_origin(&para_id).unwrap();
 
         Pallet::<T>::start_assignment(
diff --git a/pallets/data-preservers/src/lib.rs b/pallets/data-preservers/src/lib.rs
index d1b9cf6be..2a89e4bd9 100644
--- a/pallets/data-preservers/src/lib.rs
+++ b/pallets/data-preservers/src/lib.rs
@@ -33,6 +33,9 @@ mod tests;
 #[cfg(any(test, feature = "runtime-benchmarks"))]
 mod benchmarks;
 
+#[cfg(feature = "runtime-benchmarks")]
+pub use benchmarks::ArgumentFactory;
+
 pub mod weights;
 pub use weights::WeightInfo;
 
@@ -172,6 +175,10 @@ pub mod pallet {
         type AssignmentPayment: AssignmentPayment<Self::AccountId>;
 
         type WeightInfo: WeightInfo;
+
+        /// Helper type for benchmarks.
+        #[cfg(feature = "runtime-benchmarks")]
+        type BenchmarkHelper: ArgumentFactory<ParaId>;
     }
 
     #[pallet::event]
diff --git a/runtime/dancebox/src/lib.rs b/runtime/dancebox/src/lib.rs
index 80217ec8c..08d1feaec 100644
--- a/runtime/dancebox/src/lib.rs
+++ b/runtime/dancebox/src/lib.rs
@@ -1165,6 +1165,8 @@ impl pallet_data_preservers::Config for Runtime {
     type MaxAssignmentsPerParaId = MaxAssignmentsPerParaId;
     type MaxNodeUrlLen = MaxNodeUrlLen;
     type MaxParaIdsVecLen = MaxLengthParaIds;
+    #[cfg(feature = "runtime-benchmarks")]
+    type BenchmarkHelper = ();
 }
 
 impl pallet_author_noting::Config for Runtime {
diff --git a/runtime/flashbox/src/lib.rs b/runtime/flashbox/src/lib.rs
index 17678bd29..98189fdff 100644
--- a/runtime/flashbox/src/lib.rs
+++ b/runtime/flashbox/src/lib.rs
@@ -1013,6 +1013,9 @@ impl pallet_data_preservers::Config for Runtime {
     type MaxAssignmentsPerParaId = MaxAssignmentsPerParaId;
     type MaxNodeUrlLen = MaxNodeUrlLen;
     type MaxParaIdsVecLen = MaxLengthParaIds;
+
+    #[cfg(feature = "runtime-benchmarks")]
+    type BenchmarkHelper = ();
 }
 
 impl pallet_author_noting::Config for Runtime {
diff --git a/solo-chains/runtime/dancelight/src/lib.rs b/solo-chains/runtime/dancelight/src/lib.rs
index 4e211b020..7bf87227f 100644
--- a/solo-chains/runtime/dancelight/src/lib.rs
+++ b/solo-chains/runtime/dancelight/src/lib.rs
@@ -1626,6 +1626,8 @@ impl pallet_data_preservers::Config for Runtime {
     type MaxAssignmentsPerParaId = MaxAssignmentsPerParaId;
     type MaxNodeUrlLen = MaxNodeUrlLen;
     type MaxParaIdsVecLen = MaxLengthParaIds;
+    #[cfg(feature = "runtime-benchmarks")]
+    type BenchmarkHelper = ();
 }
 
 parameter_types! {
@@ -2045,9 +2047,7 @@ impl pallet_registrar::Config for Runtime {
     type DepositAmount = DepositAmount;
     type RegistrarHooks = DancelightRegistrarHooks;
     type RuntimeHoldReason = RuntimeHoldReason;
-    #[cfg(feature = "runtime-benchmarks")]
-    type InnerRegistrar = ();
-    #[cfg(not(feature = "runtime-benchmarks"))]
+
     type InnerRegistrar = InnerDancelightRegistrar<
         Runtime,
         AccountId,