From 9cdbdc5a73f3c2d1d9afdf266799f249c3c04c2d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20K=C3=B6cher?= <git@kchr.de>
Date: Tue, 17 Sep 2024 15:21:50 +0200
Subject: [PATCH] pallet-treasury: Improve `remove_approval` benchmark (#5713)

When `SpendOrigin` doesn't return any `succesful_origin`, it doesn't
mean that `RejectOrigin` will do the same. Thus, this pr fixes a
potential wrong benchmarked weight for when `SpendOrigin` is set to e.g.
`NeverOrigin`.
---
 prdoc/pr_5713.prdoc                          | 10 +++++++
 substrate/frame/treasury/src/benchmarking.rs | 29 +++++++++++++++-----
 2 files changed, 32 insertions(+), 7 deletions(-)
 create mode 100644 prdoc/pr_5713.prdoc

diff --git a/prdoc/pr_5713.prdoc b/prdoc/pr_5713.prdoc
new file mode 100644
index 000000000000..54d3619cdcaf
--- /dev/null
+++ b/prdoc/pr_5713.prdoc
@@ -0,0 +1,10 @@
+title: "pallet-treasury: Improve `remove_approval` benchmark"
+
+doc:
+  - audience: Runtime Dev
+    description: |
+      Fix the `remove_approval` benchmark when `SpendOrigin` doesn't return any `succesful_origin`.
+
+crates:
+  - name: pallet-treasury
+    bump: patch
diff --git a/substrate/frame/treasury/src/benchmarking.rs b/substrate/frame/treasury/src/benchmarking.rs
index 0bac78503f41..650e5376fa4b 100644
--- a/substrate/frame/treasury/src/benchmarking.rs
+++ b/substrate/frame/treasury/src/benchmarking.rs
@@ -133,16 +133,31 @@ mod benchmarks {
 
 	#[benchmark]
 	fn remove_approval() -> Result<(), BenchmarkError> {
-		let origin =
-			T::SpendOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
-		let (_, value, beneficiary_lookup) = setup_proposal::<T, _>(SEED);
-		Treasury::<T, _>::spend_local(origin, value, beneficiary_lookup)?;
-		let proposal_id = ProposalCount::<T, _>::get() - 1;
+		let (spend_exists, proposal_id) =
+			if let Ok(origin) = T::SpendOrigin::try_successful_origin() {
+				let (_, value, beneficiary_lookup) = setup_proposal::<T, _>(SEED);
+				Treasury::<T, _>::spend_local(origin, value, beneficiary_lookup)?;
+				let proposal_id = ProposalCount::<T, _>::get() - 1;
+
+				(true, proposal_id)
+			} else {
+				(false, 0)
+			};
+
 		let reject_origin =
 			T::RejectOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
 
-		#[extrinsic_call]
-		_(reject_origin as T::RuntimeOrigin, proposal_id);
+		#[block]
+		{
+			let res =
+				Treasury::<T, _>::remove_approval(reject_origin as T::RuntimeOrigin, proposal_id);
+
+			if spend_exists {
+				assert_ok!(res);
+			} else {
+				assert_err!(res, Error::<T, _>::ProposalNotApproved);
+			}
+		}
 
 		Ok(())
 	}