1use super::*;
21use crate::Pallet as AssetRewards;
22use frame_benchmarking::{v2::*, whitelisted_caller, BenchmarkError};
23use frame_support::{
24 assert_ok,
25 traits::{
26 fungibles::{Create, Inspect, Mutate},
27 Consideration, EnsureOrigin, Footprint,
28 },
29};
30use frame_system::{Pallet as System, RawOrigin};
31use sp_runtime::{traits::One, Saturating};
32use sp_std::prelude::*;
33
34pub trait BenchmarkHelper<AssetId> {
36 fn staked_asset() -> AssetId;
40 fn reward_asset() -> AssetId;
44}
45
46fn pool_expire<T: Config>() -> DispatchTime<BlockNumberFor<T>> {
47 DispatchTime::At(BlockNumberFor::<T>::from(100u32))
48}
49
50fn create_reward_pool<T: Config>() -> Result<T::RuntimeOrigin, BenchmarkError>
51where
52 T::Assets: Create<T::AccountId> + Mutate<T::AccountId>,
53{
54 let caller_origin =
55 T::CreatePoolOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
56 let caller = T::CreatePoolOrigin::ensure_origin(caller_origin.clone()).unwrap();
57
58 let footprint = Footprint::from_mel::<(PoolId, PoolInfoFor<T>)>();
59 T::Consideration::ensure_successful(&caller, footprint);
60
61 let staked_asset = T::BenchmarkHelper::staked_asset();
62 let reward_asset = T::BenchmarkHelper::reward_asset();
63
64 let min_staked_balance =
65 T::Assets::minimum_balance(staked_asset.clone()).max(T::Balance::one());
66 if !T::Assets::asset_exists(staked_asset.clone()) {
67 assert_ok!(T::Assets::create(
68 staked_asset.clone(),
69 caller.clone(),
70 true,
71 min_staked_balance
72 ));
73 }
74 let min_reward_balance =
75 T::Assets::minimum_balance(reward_asset.clone()).max(T::Balance::one());
76 if !T::Assets::asset_exists(reward_asset.clone()) {
77 assert_ok!(T::Assets::create(
78 reward_asset.clone(),
79 caller.clone(),
80 true,
81 min_reward_balance
82 ));
83 }
84
85 assert_ok!(AssetRewards::<T>::create_pool(
86 caller_origin.clone(),
87 Box::new(staked_asset),
88 Box::new(reward_asset),
89 min_reward_balance,
91 pool_expire::<T>(),
92 Some(caller),
93 ));
94
95 Ok(caller_origin)
96}
97
98fn mint_into<T: Config>(caller: &T::AccountId, asset: &T::AssetId) -> T::Balance
99where
100 T::Assets: Mutate<T::AccountId>,
101{
102 let min_balance = T::Assets::minimum_balance(asset.clone());
103 assert_ok!(T::Assets::mint_into(
104 asset.clone(),
105 &caller,
106 min_balance.saturating_mul(10u32.into())
107 ));
108 min_balance
109}
110
111fn assert_last_event<T: Config>(generic_event: <T as Config>::RuntimeEvent) {
112 System::<T>::assert_last_event(generic_event.into());
113}
114
115#[benchmarks(where T::Assets: Create<T::AccountId> + Mutate<T::AccountId>)]
116mod benchmarks {
117 use super::*;
118
119 #[benchmark]
120 fn create_pool() -> Result<(), BenchmarkError> {
121 let caller_origin =
122 T::CreatePoolOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
123 let caller = T::CreatePoolOrigin::ensure_origin(caller_origin.clone()).unwrap();
124
125 let footprint = Footprint::from_mel::<(PoolId, PoolInfoFor<T>)>();
126 T::Consideration::ensure_successful(&caller, footprint);
127
128 let staked_asset = T::BenchmarkHelper::staked_asset();
129 let reward_asset = T::BenchmarkHelper::reward_asset();
130
131 let min_balance = T::Assets::minimum_balance(staked_asset.clone()).max(T::Balance::one());
132 if !T::Assets::asset_exists(staked_asset.clone()) {
133 assert_ok!(T::Assets::create(staked_asset.clone(), caller.clone(), true, min_balance));
134 }
135 let min_balance = T::Assets::minimum_balance(reward_asset.clone()).max(T::Balance::one());
136 if !T::Assets::asset_exists(reward_asset.clone()) {
137 assert_ok!(T::Assets::create(reward_asset.clone(), caller.clone(), true, min_balance));
138 }
139
140 #[extrinsic_call]
141 _(
142 caller_origin as T::RuntimeOrigin,
143 Box::new(staked_asset.clone()),
144 Box::new(reward_asset.clone()),
145 min_balance,
146 pool_expire::<T>(),
147 Some(caller.clone()),
148 );
149
150 assert_last_event::<T>(
151 Event::PoolCreated {
152 creator: caller.clone(),
153 admin: caller,
154 staked_asset_id: staked_asset,
155 reward_asset_id: reward_asset,
156 reward_rate_per_block: min_balance,
157 expiry_block: pool_expire::<T>()
158 .evaluate(T::BlockNumberProvider::current_block_number()),
159 pool_id: 0,
160 }
161 .into(),
162 );
163
164 Ok(())
165 }
166
167 #[benchmark]
168 fn stake() -> Result<(), BenchmarkError> {
169 create_reward_pool::<T>()?;
170
171 let staker: T::AccountId = whitelisted_caller();
172 let min_balance = mint_into::<T>(&staker, &T::BenchmarkHelper::staked_asset());
173
174 assert_ok!(AssetRewards::<T>::stake(
176 RawOrigin::Signed(staker.clone()).into(),
177 0,
178 min_balance
179 ));
180
181 #[extrinsic_call]
182 _(RawOrigin::Signed(staker.clone()), 0, min_balance);
183
184 assert_last_event::<T>(Event::Staked { staker, pool_id: 0, amount: min_balance }.into());
185
186 Ok(())
187 }
188
189 #[benchmark]
190 fn unstake() -> Result<(), BenchmarkError> {
191 create_reward_pool::<T>()?;
192
193 let staker: T::AccountId = whitelisted_caller();
194 let min_balance = mint_into::<T>(&staker, &T::BenchmarkHelper::staked_asset());
195
196 assert_ok!(AssetRewards::<T>::stake(
197 RawOrigin::Signed(staker.clone()).into(),
198 0,
199 min_balance,
200 ));
201
202 #[extrinsic_call]
203 _(RawOrigin::Signed(staker.clone()), 0, min_balance, None);
204
205 assert_last_event::<T>(
206 Event::Unstaked { caller: staker.clone(), staker, pool_id: 0, amount: min_balance }
207 .into(),
208 );
209
210 Ok(())
211 }
212
213 #[benchmark]
214 fn harvest_rewards() -> Result<(), BenchmarkError> {
215 create_reward_pool::<T>()?;
216
217 let pool_acc = AssetRewards::<T>::pool_account_id(&0u32);
218 let min_reward_balance = mint_into::<T>(&pool_acc, &T::BenchmarkHelper::reward_asset());
219
220 let staker = whitelisted_caller();
221 let _ = mint_into::<T>(&staker, &T::BenchmarkHelper::staked_asset());
222 assert_ok!(AssetRewards::<T>::stake(
223 RawOrigin::Signed(staker.clone()).into(),
224 0,
225 T::Balance::one(),
226 ));
227
228 T::BlockNumberProvider::set_block_number(
229 T::BlockNumberProvider::current_block_number() + BlockNumberFor::<T>::one(),
230 );
231
232 #[extrinsic_call]
233 _(RawOrigin::Signed(staker.clone()), 0, None);
234
235 assert_last_event::<T>(
236 Event::RewardsHarvested {
237 caller: staker.clone(),
238 staker,
239 pool_id: 0,
240 amount: min_reward_balance,
241 }
242 .into(),
243 );
244
245 Ok(())
246 }
247
248 #[benchmark]
249 fn set_pool_reward_rate_per_block() -> Result<(), BenchmarkError> {
250 let caller_origin = create_reward_pool::<T>()?;
251
252 {
254 let staker: T::AccountId = whitelisted_caller();
255 let min_balance = mint_into::<T>(&staker, &T::BenchmarkHelper::staked_asset());
256
257 assert_ok!(AssetRewards::<T>::stake(RawOrigin::Signed(staker).into(), 0, min_balance));
258 }
259
260 let new_reward_rate_per_block =
261 T::Assets::minimum_balance(T::BenchmarkHelper::reward_asset()).max(T::Balance::one()) +
262 T::Balance::one();
263
264 #[extrinsic_call]
265 _(caller_origin as T::RuntimeOrigin, 0, new_reward_rate_per_block);
266
267 assert_last_event::<T>(
268 Event::PoolRewardRateModified { pool_id: 0, new_reward_rate_per_block }.into(),
269 );
270 Ok(())
271 }
272
273 #[benchmark]
274 fn set_pool_admin() -> Result<(), BenchmarkError> {
275 let caller_origin = create_reward_pool::<T>()?;
276 let new_admin: T::AccountId = whitelisted_caller();
277
278 #[extrinsic_call]
279 _(caller_origin as T::RuntimeOrigin, 0, new_admin.clone());
280
281 assert_last_event::<T>(Event::PoolAdminModified { pool_id: 0, new_admin }.into());
282
283 Ok(())
284 }
285
286 #[benchmark]
287 fn set_pool_expiry_block() -> Result<(), BenchmarkError> {
288 let create_origin = create_reward_pool::<T>()?;
289
290 {
292 let staker: T::AccountId = whitelisted_caller();
293 let min_balance = mint_into::<T>(&staker, &T::BenchmarkHelper::staked_asset());
294
295 assert_ok!(AssetRewards::<T>::stake(RawOrigin::Signed(staker).into(), 0, min_balance));
296 }
297
298 let new_expiry_block = pool_expire::<T>()
299 .evaluate(T::BlockNumberProvider::current_block_number()) +
300 BlockNumberFor::<T>::one();
301
302 #[extrinsic_call]
303 _(create_origin as T::RuntimeOrigin, 0, DispatchTime::At(new_expiry_block));
304
305 assert_last_event::<T>(
306 Event::PoolExpiryBlockModified { pool_id: 0, new_expiry_block }.into(),
307 );
308
309 Ok(())
310 }
311
312 #[benchmark]
313 fn deposit_reward_tokens() -> Result<(), BenchmarkError> {
314 create_reward_pool::<T>()?;
315 let caller = whitelisted_caller();
316
317 let reward_asset = T::BenchmarkHelper::reward_asset();
318 let pool_acc = AssetRewards::<T>::pool_account_id(&0u32);
319 let min_balance = mint_into::<T>(&caller, &reward_asset);
320
321 let balance_before = T::Assets::balance(reward_asset.clone(), &pool_acc);
322
323 #[extrinsic_call]
324 _(RawOrigin::Signed(caller), 0, min_balance);
325
326 let balance_after = T::Assets::balance(reward_asset, &pool_acc);
327
328 assert_eq!(balance_after, balance_before + min_balance);
329
330 Ok(())
331 }
332
333 #[benchmark]
334 fn cleanup_pool() -> Result<(), BenchmarkError> {
335 let create_origin = create_reward_pool::<T>()?;
336 let caller = T::CreatePoolOrigin::ensure_origin(create_origin.clone()).unwrap();
337
338 {
340 let caller = whitelisted_caller();
341 let reward_asset = T::BenchmarkHelper::reward_asset();
342 let min_balance = mint_into::<T>(&caller, &reward_asset);
343 assert_ok!(AssetRewards::<T>::deposit_reward_tokens(
344 RawOrigin::Signed(caller).into(),
345 0,
346 min_balance
347 ));
348 }
349
350 #[extrinsic_call]
351 _(RawOrigin::Signed(caller), 0);
352
353 assert_last_event::<T>(Event::PoolCleanedUp { pool_id: 0 }.into());
354
355 Ok(())
356 }
357
358 impl_benchmark_test_suite!(AssetRewards, crate::mock::new_test_ext(), crate::mock::MockRuntime);
359}