From b41419ae42c083b536fb23ce44f7cc8a8a1fd9f6 Mon Sep 17 00:00:00 2001
From: Xavier Lau <xavier@inv.cafe>
Date: Tue, 29 Oct 2024 21:29:15 +0800
Subject: [PATCH] Migrate pallet-timestamp benchmark to v2 (#6258)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Part of:

- #6202

---------

Co-authored-by: Dónal Murray <donalm@seadanda.dev>
---
 substrate/frame/timestamp/src/benchmarking.rs | 60 ++++++++++++-------
 1 file changed, 40 insertions(+), 20 deletions(-)

diff --git a/substrate/frame/timestamp/src/benchmarking.rs b/substrate/frame/timestamp/src/benchmarking.rs
index d8c27b4967af..1ba8c4195de9 100644
--- a/substrate/frame/timestamp/src/benchmarking.rs
+++ b/substrate/frame/timestamp/src/benchmarking.rs
@@ -19,43 +19,63 @@
 
 #![cfg(feature = "runtime-benchmarks")]
 
-use super::*;
-use frame_benchmarking::v1::benchmarks;
-use frame_support::{ensure, traits::OnFinalize};
+use frame_benchmarking::{benchmarking::add_to_whitelist, v2::*};
+use frame_support::traits::OnFinalize;
 use frame_system::RawOrigin;
 use sp_storage::TrackedStorageKey;
 
-use crate::{Now, Pallet as Timestamp};
+use crate::*;
 
 const MAX_TIME: u32 = 100;
 
-benchmarks! {
-	set {
+#[benchmarks]
+mod benchmarks {
+	use super::*;
+
+	#[benchmark]
+	fn set() {
 		let t = MAX_TIME;
 		// Ignore write to `DidUpdate` since it transient.
-		let did_update_key = crate::DidUpdate::<T>::hashed_key().to_vec();
-		frame_benchmarking::benchmarking::add_to_whitelist(TrackedStorageKey {
+		let did_update_key = DidUpdate::<T>::hashed_key().to_vec();
+		add_to_whitelist(TrackedStorageKey {
 			key: did_update_key,
 			reads: 0,
 			writes: 1,
 			whitelisted: false,
 		});
-	}: _(RawOrigin::None, t.into())
-	verify {
-		ensure!(Now::<T>::get() == t.into(), "Time was not set.");
+
+		#[extrinsic_call]
+		_(RawOrigin::None, t.into());
+
+		assert_eq!(Now::<T>::get(), t.into(), "Time was not set.");
 	}
 
-	on_finalize {
+	#[benchmark]
+	fn on_finalize() {
 		let t = MAX_TIME;
-		Timestamp::<T>::set(RawOrigin::None.into(), t.into())?;
-		ensure!(DidUpdate::<T>::exists(), "Time was not set.");
+		Pallet::<T>::set(RawOrigin::None.into(), t.into()).unwrap();
+		assert!(DidUpdate::<T>::exists(), "Time was not set.");
+
 		// Ignore read/write to `DidUpdate` since it is transient.
-		let did_update_key = crate::DidUpdate::<T>::hashed_key().to_vec();
-		frame_benchmarking::benchmarking::add_to_whitelist(did_update_key.into());
-	}: { Timestamp::<T>::on_finalize(t.into()); }
-	verify {
-		ensure!(!DidUpdate::<T>::exists(), "Time was not removed.");
+		let did_update_key = DidUpdate::<T>::hashed_key().to_vec();
+		add_to_whitelist(TrackedStorageKey {
+			key: did_update_key,
+			reads: 0,
+			writes: 1,
+			whitelisted: false,
+		});
+
+		#[block]
+		{
+			Pallet::<T>::on_finalize(t.into());
+		}
+
+		assert!(!DidUpdate::<T>::exists(), "Time was not removed.");
 	}
 
-	impl_benchmark_test_suite!(Timestamp, crate::mock::new_test_ext(), crate::mock::Test);
+	impl_benchmark_test_suite! {
+		Pallet,
+		mock::new_test_ext(),
+		mock::Test
+	}
 }