From 500497cfa0f0f828e6e97105f41f15102b72df8e Mon Sep 17 00:00:00 2001
From: bgallois <benjamin@gallois.cc>
Date: Tue, 24 Oct 2023 19:25:22 +0200
Subject: [PATCH] fix(pallet_identity): fix benchmarks

---
 pallets/distance/src/lib.rs          | 21 +++++++++++++++++++++
 pallets/identity/src/benchmarking.rs |  1 +
 pallets/identity/src/lib.rs          |  3 +++
 pallets/identity/src/mock.rs         |  2 ++
 pallets/identity/src/traits.rs       | 10 ++++++++++
 runtime/common/src/pallets_config.rs |  2 ++
 runtime/common/src/providers.rs      | 21 +++++++++++++++++++++
 7 files changed, 60 insertions(+)

diff --git a/pallets/distance/src/lib.rs b/pallets/distance/src/lib.rs
index 3ed6ab290..92ed6acc0 100644
--- a/pallets/distance/src/lib.rs
+++ b/pallets/distance/src/lib.rs
@@ -276,6 +276,27 @@ pub mod pallet {
         }
     }
 
+    // BENCHMARK FUNCTIONS //
+
+    impl<T: Config> Pallet<T> {
+        /// Force the distance status using IdtyIndex and AccountId
+        /// only to prepare identity for benchmarking.
+        pub fn set_distance_status(
+            identity: <T as pallet_identity::Config>::IdtyIndex,
+            status: Option<(<T as frame_system::Config>::AccountId, DistanceStatus)>,
+        ) -> DispatchResult {
+            IdentityDistanceStatus::<T>::set(identity, status);
+            DistanceStatusExpireOn::<T>::mutate(
+                pallet_session::CurrentIndex::<T>::get() + T::ResultExpiration::get(),
+                move |identities| {
+                    identities
+                        .try_push(identity)
+                        .map_err(|_| Error::<T>::ManyEvaluationsInBlock.into())
+                },
+            )
+        }
+    }
+
     // INTERNAL FUNCTIONS //
 
     impl<T: Config> Pallet<T> {
diff --git a/pallets/identity/src/benchmarking.rs b/pallets/identity/src/benchmarking.rs
index 1cee79c28..27894c2aa 100644
--- a/pallets/identity/src/benchmarking.rs
+++ b/pallets/identity/src/benchmarking.rs
@@ -147,6 +147,7 @@ benchmarks! {
         let owner_key: T::AccountId = Identities::<T>::get(index).unwrap().owner_key;
         let owner_key_origin: <T as frame_system::Config>::RuntimeOrigin = RawOrigin::Signed(owner_key.clone()).into();
         Pallet::<T>::confirm_identity(owner_key_origin.clone(), name.clone())?;
+        T::BenchmarkSetupHandler::force_status_ok(&index, &owner_key);
     }: _<T::RuntimeOrigin>(caller_origin, index)
     verify {
         assert_has_event::<T>(Event::<T>::IdtyValidated { idty_index: index }.into());
diff --git a/pallets/identity/src/lib.rs b/pallets/identity/src/lib.rs
index 25733edb5..46b953a07 100644
--- a/pallets/identity/src/lib.rs
+++ b/pallets/identity/src/lib.rs
@@ -114,6 +114,9 @@ pub mod pallet {
         type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;
         /// Type representing the weight of this pallet
         type WeightInfo: WeightInfo;
+        /// Type representing the a distance handler to prepare identity for benchmarking
+        #[cfg(feature = "runtime-benchmarks")]
+        type BenchmarkSetupHandler: SetDistance<Self::IdtyIndex, Self::AccountId>;
     }
 
     // GENESIS STUFF //
diff --git a/pallets/identity/src/mock.rs b/pallets/identity/src/mock.rs
index ce6032d08..7621c4955 100644
--- a/pallets/identity/src/mock.rs
+++ b/pallets/identity/src/mock.rs
@@ -117,6 +117,8 @@ impl pallet_identity::Config for Test {
     type RevocationSignature = Signature;
     type RuntimeEvent = RuntimeEvent;
     type WeightInfo = ();
+    #[cfg(feature = "runtime-benchmarks")]
+    type BenchmarkSetupHandler = ();
 }
 
 // Build genesis storage according to the mock runtime.
diff --git a/pallets/identity/src/traits.rs b/pallets/identity/src/traits.rs
index 54b50a06c..1d8dd0152 100644
--- a/pallets/identity/src/traits.rs
+++ b/pallets/identity/src/traits.rs
@@ -76,3 +76,13 @@ impl<IndtyIndex> RemoveIdentityConsumers<IndtyIndex> for () {
         Weight::zero()
     }
 }
+
+#[cfg(feature = "runtime-benchmarks")]
+pub trait SetDistance<IndtyIndex, AccountId> {
+    fn force_status_ok(idty_index: &IndtyIndex, account: &AccountId) -> ();
+}
+
+#[cfg(feature = "runtime-benchmarks")]
+impl<IdtyIndex, AccountId> SetDistance<IdtyIndex, AccountId> for () {
+    fn force_status_ok(_idty_id: &IdtyIndex, _account: &AccountId) -> () {}
+}
diff --git a/runtime/common/src/pallets_config.rs b/runtime/common/src/pallets_config.rs
index 665bd8193..a2bb56c9c 100644
--- a/runtime/common/src/pallets_config.rs
+++ b/runtime/common/src/pallets_config.rs
@@ -454,6 +454,8 @@ parameter_types! {
             type RevocationSignature = Signature;
             type RuntimeEvent = RuntimeEvent;
             type WeightInfo = common_runtime::weights::pallet_identity::WeightInfo<Runtime>;
+            #[cfg(feature = "runtime-benchmarks")]
+            type BenchmarkSetupHandler = common_runtime::providers::BenchmarkSetupHandler<Runtime>;
         }
 
         impl pallet_membership::Config<frame_support::instances::Instance1> for Runtime {
diff --git a/runtime/common/src/providers.rs b/runtime/common/src/providers.rs
index 151f70ece..2079dfc3e 100644
--- a/runtime/common/src/providers.rs
+++ b/runtime/common/src/providers.rs
@@ -121,3 +121,24 @@ where
         )
     }
 }
+
+#[cfg(feature = "runtime-benchmarks")]
+pub struct BenchmarkSetupHandler<T>(PhantomData<T>);
+
+#[cfg(feature = "runtime-benchmarks")]
+impl<T>
+    pallet_identity::traits::SetDistance<<T as pallet_identity::Config>::IdtyIndex, T::AccountId>
+    for BenchmarkSetupHandler<T>
+where
+    T: pallet_distance::Config,
+{
+    fn force_status_ok(
+        idty_id: &<T as pallet_identity::Config>::IdtyIndex,
+        account: &<T as frame_system::Config>::AccountId,
+    ) -> () {
+        let _ = pallet_distance::Pallet::<T>::set_distance_status(
+            *idty_id,
+            Some((account.clone(), pallet_distance::DistanceStatus::Valid)),
+        );
+    }
+}
-- 
GitLab