1#![cfg(feature = "runtime-benchmarks")]
21
22use super::{Pallet as Treasury, *};
23
24use frame_benchmarking::{
25 v1::{account, BenchmarkError},
26 v2::*,
27};
28use frame_support::{
29 assert_err, assert_ok, ensure,
30 traits::{
31 tokens::{ConversionFromAssetBalance, PaymentStatus},
32 EnsureOrigin, OnInitialize,
33 },
34};
35use frame_system::RawOrigin;
36use sp_core::crypto::FromEntropy;
37
38pub trait ArgumentsFactory<AssetKind, Beneficiary> {
40 fn create_asset_kind(seed: u32) -> AssetKind;
42 fn create_beneficiary(seed: [u8; 32]) -> Beneficiary;
44}
45
46impl<AssetKind, Beneficiary> ArgumentsFactory<AssetKind, Beneficiary> for ()
48where
49 AssetKind: FromEntropy,
50 Beneficiary: FromEntropy,
51{
52 fn create_asset_kind(seed: u32) -> AssetKind {
53 AssetKind::from_entropy(&mut seed.encode().as_slice()).unwrap()
54 }
55 fn create_beneficiary(seed: [u8; 32]) -> Beneficiary {
56 Beneficiary::from_entropy(&mut seed.as_slice()).unwrap()
57 }
58}
59
60const SEED: u32 = 0;
61
62fn setup_proposal<T: Config<I>, I: 'static>(
64 u: u32,
65) -> (T::AccountId, BalanceOf<T, I>, AccountIdLookupOf<T>) {
66 let caller = account("caller", u, SEED);
67 let value: BalanceOf<T, I> = T::Currency::minimum_balance() * 100u32.into();
68 let _ = T::Currency::make_free_balance_be(&caller, value);
69 let beneficiary = account("beneficiary", u, SEED);
70 let beneficiary_lookup = T::Lookup::unlookup(beneficiary);
71 (caller, value, beneficiary_lookup)
72}
73
74fn create_approved_proposals<T: Config<I>, I: 'static>(n: u32) -> Result<(), &'static str> {
76 let spender = T::SpendOrigin::try_successful_origin();
77
78 for i in 0..n {
79 let (_, value, lookup) = setup_proposal::<T, I>(i);
80
81 #[allow(deprecated)]
82 if let Ok(origin) = &spender {
83 Treasury::<T, I>::spend_local(origin.clone(), value, lookup)?;
84 }
85 }
86
87 if spender.is_ok() {
88 ensure!(Approvals::<T, I>::get().len() == n as usize, "Not all approved");
89 }
90 Ok(())
91}
92
93fn setup_pot_account<T: Config<I>, I: 'static>() {
94 let pot_account = Treasury::<T, I>::account_id();
95 let value = T::Currency::minimum_balance().saturating_mul(1_000_000_000u32.into());
96 let _ = T::Currency::make_free_balance_be(&pot_account, value);
97}
98
99fn assert_last_event<T: Config<I>, I: 'static>(generic_event: <T as Config<I>>::RuntimeEvent) {
100 frame_system::Pallet::<T>::assert_last_event(generic_event.into());
101}
102
103fn create_spend_arguments<T: Config<I>, I: 'static>(
105 seed: u32,
106) -> (T::AssetKind, AssetBalanceOf<T, I>, T::Beneficiary, BeneficiaryLookupOf<T, I>) {
107 let asset_kind = T::BenchmarkHelper::create_asset_kind(seed);
108 let beneficiary = T::BenchmarkHelper::create_beneficiary([seed.try_into().unwrap(); 32]);
109 let beneficiary_lookup = T::BeneficiaryLookup::unlookup(beneficiary.clone());
110 (asset_kind, 100u32.into(), beneficiary, beneficiary_lookup)
111}
112
113#[instance_benchmarks]
114mod benchmarks {
115 use super::*;
116
117 #[benchmark]
120 fn spend_local() -> Result<(), BenchmarkError> {
121 let (_, value, beneficiary_lookup) = setup_proposal::<T, _>(SEED);
122 let origin =
123 T::SpendOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
124 let beneficiary = T::Lookup::lookup(beneficiary_lookup.clone()).unwrap();
125
126 #[extrinsic_call]
127 _(origin as T::RuntimeOrigin, value, beneficiary_lookup);
128
129 assert_last_event::<T, I>(
130 Event::SpendApproved { proposal_index: 0, amount: value, beneficiary }.into(),
131 );
132 Ok(())
133 }
134
135 #[benchmark]
136 fn remove_approval() -> Result<(), BenchmarkError> {
137 let (spend_exists, proposal_id) =
138 if let Ok(origin) = T::SpendOrigin::try_successful_origin() {
139 let (_, value, beneficiary_lookup) = setup_proposal::<T, _>(SEED);
140 #[allow(deprecated)]
141 Treasury::<T, _>::spend_local(origin, value, beneficiary_lookup)?;
142 let proposal_id = ProposalCount::<T, _>::get() - 1;
143
144 (true, proposal_id)
145 } else {
146 (false, 0)
147 };
148
149 let reject_origin =
150 T::RejectOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
151
152 #[block]
153 {
154 #[allow(deprecated)]
155 let res = Treasury::<T, _>::remove_approval(reject_origin as T::RuntimeOrigin, proposal_id);
156
157 if spend_exists {
158 assert_ok!(res);
159 } else {
160 assert_err!(res, Error::<T, _>::ProposalNotApproved);
161 }
162 }
163
164 Ok(())
165 }
166
167 #[benchmark]
168 fn on_initialize_proposals(
169 p: Linear<0, { T::MaxApprovals::get() - 1 }>,
170 ) -> Result<(), BenchmarkError> {
171 setup_pot_account::<T, _>();
172 create_approved_proposals::<T, _>(p)?;
173
174 #[block]
175 {
176 Treasury::<T, _>::on_initialize(0u32.into());
177 }
178
179 Ok(())
180 }
181
182 #[benchmark]
185 fn spend() -> Result<(), BenchmarkError> {
186 let origin =
187 T::SpendOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
188 let (asset_kind, amount, beneficiary, beneficiary_lookup) =
189 create_spend_arguments::<T, _>(SEED);
190 T::BalanceConverter::ensure_successful(asset_kind.clone());
191
192 #[extrinsic_call]
193 _(
194 origin as T::RuntimeOrigin,
195 Box::new(asset_kind.clone()),
196 amount,
197 Box::new(beneficiary_lookup),
198 None,
199 );
200
201 let valid_from = T::BlockNumberProvider::current_block_number();
202 let expire_at = valid_from.saturating_add(T::PayoutPeriod::get());
203 assert_last_event::<T, I>(
204 Event::AssetSpendApproved {
205 index: 0,
206 asset_kind,
207 amount,
208 beneficiary,
209 valid_from,
210 expire_at,
211 }
212 .into(),
213 );
214 Ok(())
215 }
216
217 #[benchmark]
218 fn payout() -> Result<(), BenchmarkError> {
219 let (asset_kind, amount, beneficiary, beneficiary_lookup) =
220 create_spend_arguments::<T, _>(SEED);
221 T::BalanceConverter::ensure_successful(asset_kind.clone());
222
223 let spend_exists = if let Ok(origin) = T::SpendOrigin::try_successful_origin() {
224 Treasury::<T, _>::spend(
225 origin,
226 Box::new(asset_kind.clone()),
227 amount,
228 Box::new(beneficiary_lookup),
229 None,
230 )?;
231
232 true
233 } else {
234 false
235 };
236
237 T::Paymaster::ensure_successful(&beneficiary, asset_kind, amount);
238 let caller: T::AccountId = account("caller", 0, SEED);
239
240 #[block]
241 {
242 let res = Treasury::<T, _>::payout(RawOrigin::Signed(caller.clone()).into(), 0u32);
243
244 if spend_exists {
245 assert_ok!(res);
246 } else {
247 assert_err!(res, crate::Error::<T, _>::InvalidIndex);
248 }
249 }
250
251 if spend_exists {
252 let id = match Spends::<T, I>::get(0).unwrap().status {
253 PaymentState::Attempted { id, .. } => {
254 assert_ne!(T::Paymaster::check_payment(id), PaymentStatus::Failure);
255 id
256 },
257 _ => panic!("No payout attempt made"),
258 };
259 assert_last_event::<T, I>(Event::Paid { index: 0, payment_id: id }.into());
260 assert!(Treasury::<T, _>::payout(RawOrigin::Signed(caller).into(), 0u32).is_err());
261 }
262
263 Ok(())
264 }
265
266 #[benchmark]
267 fn check_status() -> Result<(), BenchmarkError> {
268 let (asset_kind, amount, beneficiary, beneficiary_lookup) =
269 create_spend_arguments::<T, _>(SEED);
270
271 T::BalanceConverter::ensure_successful(asset_kind.clone());
272 T::Paymaster::ensure_successful(&beneficiary, asset_kind.clone(), amount);
273 let caller: T::AccountId = account("caller", 0, SEED);
274
275 let spend_exists = if let Ok(origin) = T::SpendOrigin::try_successful_origin() {
276 Treasury::<T, _>::spend(
277 origin,
278 Box::new(asset_kind),
279 amount,
280 Box::new(beneficiary_lookup),
281 None,
282 )?;
283
284 Treasury::<T, _>::payout(RawOrigin::Signed(caller.clone()).into(), 0u32)?;
285 match Spends::<T, I>::get(0).unwrap().status {
286 PaymentState::Attempted { id, .. } => {
287 T::Paymaster::ensure_concluded(id);
288 },
289 _ => panic!("No payout attempt made"),
290 };
291
292 true
293 } else {
294 false
295 };
296
297 #[block]
298 {
299 let res =
300 Treasury::<T, _>::check_status(RawOrigin::Signed(caller.clone()).into(), 0u32);
301
302 if spend_exists {
303 assert_ok!(res);
304 } else {
305 assert_err!(res, crate::Error::<T, _>::InvalidIndex);
306 }
307 }
308
309 if let Some(s) = Spends::<T, I>::get(0) {
310 assert!(!matches!(s.status, PaymentState::Attempted { .. }));
311 }
312
313 Ok(())
314 }
315
316 #[benchmark]
317 fn void_spend() -> Result<(), BenchmarkError> {
318 let (asset_kind, amount, _, beneficiary_lookup) = create_spend_arguments::<T, _>(SEED);
319 T::BalanceConverter::ensure_successful(asset_kind.clone());
320 let spend_exists = if let Ok(origin) = T::SpendOrigin::try_successful_origin() {
321 Treasury::<T, _>::spend(
322 origin,
323 Box::new(asset_kind.clone()),
324 amount,
325 Box::new(beneficiary_lookup),
326 None,
327 )?;
328 assert!(Spends::<T, I>::get(0).is_some());
329
330 true
331 } else {
332 false
333 };
334
335 let origin =
336 T::RejectOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
337
338 #[block]
339 {
340 let res = Treasury::<T, _>::void_spend(origin as T::RuntimeOrigin, 0u32);
341
342 if spend_exists {
343 assert_ok!(res);
344 } else {
345 assert_err!(res, crate::Error::<T, _>::InvalidIndex);
346 }
347 }
348
349 assert!(Spends::<T, I>::get(0).is_none());
350 Ok(())
351 }
352
353 impl_benchmark_test_suite!(
354 Treasury,
355 crate::tests::ExtBuilder::default().build(),
356 crate::tests::Test
357 );
358
359 mod no_spend_origin_tests {
360 use super::*;
361
362 impl_benchmark_test_suite!(
363 Treasury,
364 crate::tests::ExtBuilder::default().spend_origin_succesful_origin_err().build(),
365 crate::tests::Test,
366 benchmarks_path = benchmarking
367 );
368 }
369}