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_prelude::BlockNumberFor, 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>().evaluate(System::<T>::block_number()),
158 pool_id: 0,
159 }
160 .into(),
161 );
162
163 Ok(())
164 }
165
166 #[benchmark]
167 fn stake() -> Result<(), BenchmarkError> {
168 create_reward_pool::<T>()?;
169
170 let staker: T::AccountId = whitelisted_caller();
171 let min_balance = mint_into::<T>(&staker, &T::BenchmarkHelper::staked_asset());
172
173 assert_ok!(AssetRewards::<T>::stake(
175 RawOrigin::Signed(staker.clone()).into(),
176 0,
177 min_balance
178 ));
179
180 #[extrinsic_call]
181 _(RawOrigin::Signed(staker.clone()), 0, min_balance);
182
183 assert_last_event::<T>(Event::Staked { staker, pool_id: 0, amount: min_balance }.into());
184
185 Ok(())
186 }
187
188 #[benchmark]
189 fn unstake() -> Result<(), BenchmarkError> {
190 create_reward_pool::<T>()?;
191
192 let staker: T::AccountId = whitelisted_caller();
193 let min_balance = mint_into::<T>(&staker, &T::BenchmarkHelper::staked_asset());
194
195 assert_ok!(AssetRewards::<T>::stake(
196 RawOrigin::Signed(staker.clone()).into(),
197 0,
198 min_balance,
199 ));
200
201 #[extrinsic_call]
202 _(RawOrigin::Signed(staker.clone()), 0, min_balance, None);
203
204 assert_last_event::<T>(
205 Event::Unstaked { caller: staker.clone(), staker, pool_id: 0, amount: min_balance }
206 .into(),
207 );
208
209 Ok(())
210 }
211
212 #[benchmark]
213 fn harvest_rewards() -> Result<(), BenchmarkError> {
214 create_reward_pool::<T>()?;
215
216 let pool_acc = AssetRewards::<T>::pool_account_id(&0u32);
217 let min_reward_balance = mint_into::<T>(&pool_acc, &T::BenchmarkHelper::reward_asset());
218
219 let staker = whitelisted_caller();
220 let _ = mint_into::<T>(&staker, &T::BenchmarkHelper::staked_asset());
221 assert_ok!(AssetRewards::<T>::stake(
222 RawOrigin::Signed(staker.clone()).into(),
223 0,
224 T::Balance::one(),
225 ));
226
227 System::<T>::set_block_number(System::<T>::block_number() + BlockNumberFor::<T>::one());
228
229 #[extrinsic_call]
230 _(RawOrigin::Signed(staker.clone()), 0, None);
231
232 assert_last_event::<T>(
233 Event::RewardsHarvested {
234 caller: staker.clone(),
235 staker,
236 pool_id: 0,
237 amount: min_reward_balance,
238 }
239 .into(),
240 );
241
242 Ok(())
243 }
244
245 #[benchmark]
246 fn set_pool_reward_rate_per_block() -> Result<(), BenchmarkError> {
247 let caller_origin = create_reward_pool::<T>()?;
248
249 {
251 let staker: T::AccountId = whitelisted_caller();
252 let min_balance = mint_into::<T>(&staker, &T::BenchmarkHelper::staked_asset());
253
254 assert_ok!(AssetRewards::<T>::stake(RawOrigin::Signed(staker).into(), 0, min_balance));
255 }
256
257 let new_reward_rate_per_block =
258 T::Assets::minimum_balance(T::BenchmarkHelper::reward_asset()).max(T::Balance::one()) +
259 T::Balance::one();
260
261 #[extrinsic_call]
262 _(caller_origin as T::RuntimeOrigin, 0, new_reward_rate_per_block);
263
264 assert_last_event::<T>(
265 Event::PoolRewardRateModified { pool_id: 0, new_reward_rate_per_block }.into(),
266 );
267 Ok(())
268 }
269
270 #[benchmark]
271 fn set_pool_admin() -> Result<(), BenchmarkError> {
272 let caller_origin = create_reward_pool::<T>()?;
273 let new_admin: T::AccountId = whitelisted_caller();
274
275 #[extrinsic_call]
276 _(caller_origin as T::RuntimeOrigin, 0, new_admin.clone());
277
278 assert_last_event::<T>(Event::PoolAdminModified { pool_id: 0, new_admin }.into());
279
280 Ok(())
281 }
282
283 #[benchmark]
284 fn set_pool_expiry_block() -> Result<(), BenchmarkError> {
285 let create_origin = create_reward_pool::<T>()?;
286
287 {
289 let staker: T::AccountId = whitelisted_caller();
290 let min_balance = mint_into::<T>(&staker, &T::BenchmarkHelper::staked_asset());
291
292 assert_ok!(AssetRewards::<T>::stake(RawOrigin::Signed(staker).into(), 0, min_balance));
293 }
294
295 let new_expiry_block =
296 pool_expire::<T>().evaluate(System::<T>::block_number()) + BlockNumberFor::<T>::one();
297
298 #[extrinsic_call]
299 _(create_origin as T::RuntimeOrigin, 0, DispatchTime::At(new_expiry_block));
300
301 assert_last_event::<T>(
302 Event::PoolExpiryBlockModified { pool_id: 0, new_expiry_block }.into(),
303 );
304
305 Ok(())
306 }
307
308 #[benchmark]
309 fn deposit_reward_tokens() -> Result<(), BenchmarkError> {
310 create_reward_pool::<T>()?;
311 let caller = whitelisted_caller();
312
313 let reward_asset = T::BenchmarkHelper::reward_asset();
314 let pool_acc = AssetRewards::<T>::pool_account_id(&0u32);
315 let min_balance = mint_into::<T>(&caller, &reward_asset);
316
317 let balance_before = T::Assets::balance(reward_asset.clone(), &pool_acc);
318
319 #[extrinsic_call]
320 _(RawOrigin::Signed(caller), 0, min_balance);
321
322 let balance_after = T::Assets::balance(reward_asset, &pool_acc);
323
324 assert_eq!(balance_after, balance_before + min_balance);
325
326 Ok(())
327 }
328
329 #[benchmark]
330 fn cleanup_pool() -> Result<(), BenchmarkError> {
331 let create_origin = create_reward_pool::<T>()?;
332 let caller = T::CreatePoolOrigin::ensure_origin(create_origin.clone()).unwrap();
333
334 {
336 let caller = whitelisted_caller();
337 let reward_asset = T::BenchmarkHelper::reward_asset();
338 let min_balance = mint_into::<T>(&caller, &reward_asset);
339 assert_ok!(AssetRewards::<T>::deposit_reward_tokens(
340 RawOrigin::Signed(caller).into(),
341 0,
342 min_balance
343 ));
344 }
345
346 #[extrinsic_call]
347 _(RawOrigin::Signed(caller), 0);
348
349 assert_last_event::<T>(Event::PoolCleanedUp { pool_id: 0 }.into());
350
351 Ok(())
352 }
353
354 impl_benchmark_test_suite!(AssetRewards, crate::mock::new_test_ext(), crate::mock::MockRuntime);
355}