1use super::*;
19use crate::{self as pools};
20use frame_support::{
21 assert_ok, derive_impl, ord_parameter_types, parameter_types,
22 traits::{fungible::Mutate, VariantCountOf},
23 PalletId,
24};
25use frame_system::{EnsureSignedBy, RawOrigin};
26use sp_runtime::{BuildStorage, DispatchResult, FixedU128};
27use sp_staking::{
28 Agent, DelegationInterface, DelegationMigrator, Delegator, OnStakingUpdate, Stake,
29};
30
31pub type BlockNumber = u64;
32pub type AccountId = u128;
33pub type Balance = u128;
34pub type RewardCounter = FixedU128;
35pub type T = Runtime;
38pub type Currency = <T as Config>::Currency;
39
40pub fn default_bonded_account() -> AccountId {
42 Pools::generate_bonded_account(1)
43}
44
45pub fn default_reward_account() -> AccountId {
47 Pools::generate_reward_account(1)
48}
49
50parameter_types! {
51 pub static MinJoinBondConfig: Balance = 2;
52 pub static CurrentEra: EraIndex = 0;
53 pub static BondingDuration: EraIndex = 3;
54 pub storage BondedBalanceMap: BTreeMap<AccountId, Balance> = Default::default();
55 pub storage UnbondingBalanceMap: BTreeMap<AccountId, Vec<(EraIndex, Balance)>> = Default::default();
57 #[derive(Clone, PartialEq)]
58 pub static MaxUnbonding: u32 = 8;
59 pub static StakingMinBond: Balance = 10;
60 pub storage Nominations: Option<Vec<AccountId>> = None;
61 pub static RestrictedAccounts: Vec<AccountId> = Vec::new();
62}
63pub struct StakingMock;
64
65impl StakingMock {
66 pub(crate) fn set_bonded_balance(who: AccountId, bonded: Balance) {
67 let mut x = BondedBalanceMap::get();
68 x.insert(who, bonded);
69 BondedBalanceMap::set(&x)
70 }
71 pub fn slash_by(pool_id: PoolId, amount: Balance) {
78 let acc = Pools::generate_bonded_account(pool_id);
79 let bonded = BondedBalanceMap::get();
80 let pre_total = bonded.get(&acc).unwrap();
81 Self::set_bonded_balance(acc, pre_total - amount);
82 DelegateMock::on_slash(acc, amount);
83 Pools::on_slash(&acc, pre_total - amount, &Default::default(), amount);
84 }
85}
86
87impl sp_staking::StakingInterface for StakingMock {
88 type Balance = Balance;
89 type AccountId = AccountId;
90 type CurrencyToVote = ();
91
92 fn minimum_nominator_bond() -> Self::Balance {
93 StakingMinBond::get()
94 }
95 fn minimum_validator_bond() -> Self::Balance {
96 StakingMinBond::get()
97 }
98
99 fn desired_validator_count() -> u32 {
100 unimplemented!("method currently not used in testing")
101 }
102
103 fn current_era() -> EraIndex {
104 CurrentEra::get()
105 }
106
107 fn bonding_duration() -> EraIndex {
108 BondingDuration::get()
109 }
110
111 fn status(
112 _: &Self::AccountId,
113 ) -> Result<sp_staking::StakerStatus<Self::AccountId>, DispatchError> {
114 Nominations::get()
115 .map(|noms| sp_staking::StakerStatus::Nominator(noms))
116 .ok_or(DispatchError::Other("NotStash"))
117 }
118
119 fn is_virtual_staker(who: &Self::AccountId) -> bool {
120 AgentBalanceMap::get().contains_key(who)
121 }
122
123 fn bond_extra(who: &Self::AccountId, extra: Self::Balance) -> DispatchResult {
124 let mut x = BondedBalanceMap::get();
125 x.get_mut(who).map(|v| *v += extra);
126 BondedBalanceMap::set(&x);
127 Ok(())
128 }
129
130 fn unbond(who: &Self::AccountId, amount: Self::Balance) -> DispatchResult {
131 let mut x = BondedBalanceMap::get();
132 *x.get_mut(who).unwrap() = x.get_mut(who).unwrap().saturating_sub(amount);
133 BondedBalanceMap::set(&x);
134
135 let era = Self::current_era();
136 let unlocking_at = era + Self::bonding_duration();
137 let mut y = UnbondingBalanceMap::get();
138 y.entry(*who).or_insert(Default::default()).push((unlocking_at, amount));
139 UnbondingBalanceMap::set(&y);
140 Ok(())
141 }
142
143 fn set_payee(_stash: &Self::AccountId, _reward_acc: &Self::AccountId) -> DispatchResult {
144 unimplemented!("method currently not used in testing")
145 }
146
147 fn chill(_: &Self::AccountId) -> sp_runtime::DispatchResult {
148 Ok(())
149 }
150
151 fn withdraw_unbonded(who: Self::AccountId, _: u32) -> Result<bool, DispatchError> {
152 let mut unbonding_map = UnbondingBalanceMap::get();
153
154 let unlocking = |pair: &Vec<(EraIndex, Balance)>| -> Balance {
156 pair.iter()
157 .try_fold(Zero::zero(), |acc: Balance, (_at, amount)| acc.checked_add(*amount))
158 .unwrap()
159 };
160
161 let staker_map = unbonding_map.get_mut(&who).ok_or("Nothing to unbond")?;
162 let unlocking_before = unlocking(&staker_map);
163
164 let current_era = Self::current_era();
165
166 staker_map.retain(|(unlocking_at, _amount)| *unlocking_at > current_era);
167
168 let withdraw_amount = unlocking_before.saturating_sub(unlocking(&staker_map));
170 Pools::on_withdraw(&who, withdraw_amount);
171 DelegateMock::on_withdraw(who, withdraw_amount);
172
173 UnbondingBalanceMap::set(&unbonding_map);
174 Ok(UnbondingBalanceMap::get().get(&who).unwrap().is_empty() &&
175 BondedBalanceMap::get().get(&who).unwrap().is_zero())
176 }
177
178 fn bond(stash: &Self::AccountId, value: Self::Balance, _: &Self::AccountId) -> DispatchResult {
179 StakingMock::set_bonded_balance(*stash, value);
180 Ok(())
181 }
182
183 fn nominate(_: &Self::AccountId, nominations: Vec<Self::AccountId>) -> DispatchResult {
184 Nominations::set(&Some(nominations));
185 Ok(())
186 }
187
188 #[cfg(feature = "runtime-benchmarks")]
189 fn nominations(_: &Self::AccountId) -> Option<Vec<Self::AccountId>> {
190 Nominations::get()
191 }
192
193 fn stash_by_ctrl(_controller: &Self::AccountId) -> Result<Self::AccountId, DispatchError> {
194 unimplemented!("method currently not used in testing")
195 }
196
197 fn stake(who: &Self::AccountId) -> Result<Stake<Balance>, DispatchError> {
198 match (UnbondingBalanceMap::get().get(who), BondedBalanceMap::get().get(who).copied()) {
199 (None, None) => Err(DispatchError::Other("balance not found")),
200 (Some(v), None) => Ok(Stake {
201 total: v.into_iter().fold(0u128, |acc, &x| acc.saturating_add(x.1)),
202 active: 0,
203 }),
204 (None, Some(v)) => Ok(Stake { total: v, active: v }),
205 (Some(a), Some(b)) => Ok(Stake {
206 total: a.into_iter().fold(0u128, |acc, &x| acc.saturating_add(x.1)) + b,
207 active: b,
208 }),
209 }
210 }
211
212 fn election_ongoing() -> bool {
213 unimplemented!("method currently not used in testing")
214 }
215
216 fn force_unstake(_who: Self::AccountId) -> sp_runtime::DispatchResult {
217 unimplemented!("method currently not used in testing")
218 }
219
220 fn is_exposed_in_era(_who: &Self::AccountId, _era: &EraIndex) -> bool {
221 unimplemented!("method currently not used in testing")
222 }
223
224 #[cfg(feature = "runtime-benchmarks")]
225 fn add_era_stakers(
226 _current_era: &EraIndex,
227 _stash: &Self::AccountId,
228 _exposures: Vec<(Self::AccountId, Self::Balance)>,
229 ) {
230 unimplemented!("method currently not used in testing")
231 }
232
233 #[cfg(any(feature = "std", feature = "runtime-benchmarks"))]
234 fn set_era(_era: EraIndex) {
235 unimplemented!("method currently not used in testing")
236 }
237
238 #[cfg(feature = "runtime-benchmarks")]
239 fn max_exposure_page_size() -> sp_staking::Page {
240 unimplemented!("method currently not used in testing")
241 }
242
243 fn slash_reward_fraction() -> Perbill {
244 unimplemented!("method currently not used in testing")
245 }
246}
247
248parameter_types! {
249 pub storage AgentBalanceMap: BTreeMap<AccountId, (Balance, Balance, Balance)> = Default::default();
251 pub storage DelegatorBalanceMap: BTreeMap<AccountId, Balance> = Default::default();
252}
253pub struct DelegateMock;
254impl DelegationInterface for DelegateMock {
255 type Balance = Balance;
256 type AccountId = AccountId;
257 fn agent_balance(agent: Agent<Self::AccountId>) -> Option<Self::Balance> {
258 AgentBalanceMap::get()
259 .get(&agent.get())
260 .copied()
261 .map(|(delegated, _, pending)| delegated - pending)
262 }
263
264 fn agent_transferable_balance(agent: Agent<Self::AccountId>) -> Option<Self::Balance> {
265 AgentBalanceMap::get()
266 .get(&agent.get())
267 .copied()
268 .map(|(_, unclaimed_withdrawals, _)| unclaimed_withdrawals)
269 }
270
271 fn delegator_balance(delegator: Delegator<Self::AccountId>) -> Option<Self::Balance> {
272 DelegatorBalanceMap::get().get(&delegator.get()).copied()
273 }
274
275 fn register_agent(
276 agent: Agent<Self::AccountId>,
277 _reward_account: &Self::AccountId,
278 ) -> DispatchResult {
279 let mut agents = AgentBalanceMap::get();
280 agents.insert(agent.get(), (0, 0, 0));
281 AgentBalanceMap::set(&agents);
282 Ok(())
283 }
284
285 fn remove_agent(agent: Agent<Self::AccountId>) -> DispatchResult {
286 let mut agents = AgentBalanceMap::get();
287 let agent = agent.get();
288 assert!(agents.contains_key(&agent));
289 agents.remove(&agent);
290 AgentBalanceMap::set(&agents);
291 Ok(())
292 }
293
294 fn delegate(
295 delegator: Delegator<Self::AccountId>,
296 agent: Agent<Self::AccountId>,
297 amount: Self::Balance,
298 ) -> DispatchResult {
299 let delegator = delegator.get();
300 let mut delegators = DelegatorBalanceMap::get();
301 delegators.entry(delegator).and_modify(|b| *b += amount).or_insert(amount);
302 DelegatorBalanceMap::set(&delegators);
303
304 let agent = agent.get();
305 let mut agents = AgentBalanceMap::get();
306 agents
307 .get_mut(&agent)
308 .map(|(d, _, _)| *d += amount)
309 .ok_or(DispatchError::Other("agent not registered"))?;
310 AgentBalanceMap::set(&agents);
311
312 if BondedBalanceMap::get().contains_key(&agent) {
313 StakingMock::bond_extra(&agent, amount)
314 } else {
315 StakingMock::bond(&agent, amount, &999)
317 }
318 }
319
320 fn withdraw_delegation(
321 delegator: Delegator<Self::AccountId>,
322 agent: Agent<Self::AccountId>,
323 amount: Self::Balance,
324 _num_slashing_spans: u32,
325 ) -> DispatchResult {
326 let mut delegators = DelegatorBalanceMap::get();
327 delegators.get_mut(&delegator.get()).map(|b| *b -= amount);
328 DelegatorBalanceMap::set(&delegators);
329
330 let mut agents = AgentBalanceMap::get();
331 agents.get_mut(&agent.get()).map(|(d, u, _)| {
332 *d -= amount;
333 *u -= amount;
334 });
335 AgentBalanceMap::set(&agents);
336
337 Ok(())
338 }
339
340 fn pending_slash(agent: Agent<Self::AccountId>) -> Option<Self::Balance> {
341 AgentBalanceMap::get()
342 .get(&agent.get())
343 .copied()
344 .map(|(_, _, pending_slash)| pending_slash)
345 }
346
347 fn delegator_slash(
348 agent: Agent<Self::AccountId>,
349 delegator: Delegator<Self::AccountId>,
350 value: Self::Balance,
351 _maybe_reporter: Option<Self::AccountId>,
352 ) -> DispatchResult {
353 let mut delegators = DelegatorBalanceMap::get();
354 delegators.get_mut(&delegator.get()).map(|b| *b -= value);
355 DelegatorBalanceMap::set(&delegators);
356
357 let mut agents = AgentBalanceMap::get();
358 agents.get_mut(&agent.get()).map(|(_, _, p)| {
359 p.saturating_reduce(value);
360 });
361 AgentBalanceMap::set(&agents);
362
363 Ok(())
364 }
365}
366
367impl DelegateMock {
368 pub fn set_agent_balance(who: AccountId, delegated: Balance) {
369 Self::set_agent_balance_full(who, delegated, 0, 0);
370 }
371
372 pub fn set_agent_balance_full(
375 who: AccountId,
376 delegated: Balance,
377 unclaimed_withdrawals: Balance,
378 pending_slash: Balance,
379 ) {
380 let mut agents = AgentBalanceMap::get();
381 agents.insert(who, (delegated, unclaimed_withdrawals, pending_slash));
382 AgentBalanceMap::set(&agents);
383 }
384
385 pub fn set_delegator_balance(who: AccountId, amount: Balance) {
386 let mut delegators = DelegatorBalanceMap::get();
387 delegators.insert(who, amount);
388 DelegatorBalanceMap::set(&delegators);
389 }
390
391 pub fn on_slash(agent: AccountId, amount: Balance) {
392 let mut agents = AgentBalanceMap::get();
393 agents.get_mut(&agent).map(|(_, _, p)| *p += amount);
394 AgentBalanceMap::set(&agents);
395 }
396
397 fn on_withdraw(agent: AccountId, amount: Balance) {
398 let mut agents = AgentBalanceMap::get();
399 agents.get_mut(&agent).map(|(_, u, _)| *u += amount);
401 AgentBalanceMap::set(&agents);
402 }
403}
404
405impl DelegationMigrator for DelegateMock {
406 type Balance = Balance;
407 type AccountId = AccountId;
408 fn migrate_nominator_to_agent(
409 _agent: Agent<Self::AccountId>,
410 _reward_account: &Self::AccountId,
411 ) -> DispatchResult {
412 unimplemented!("not used in current unit tests")
413 }
414
415 fn migrate_delegation(
416 _agent: Agent<Self::AccountId>,
417 _delegator: Delegator<Self::AccountId>,
418 _value: Self::Balance,
419 ) -> DispatchResult {
420 unimplemented!("not used in current unit tests")
421 }
422
423 #[cfg(feature = "runtime-benchmarks")]
424 fn force_kill_agent(_agent: Agent<Self::AccountId>) {
425 unimplemented!("not used in current unit tests")
426 }
427}
428
429#[derive_impl(frame_system::config_preludes::TestDefaultConfig)]
430impl frame_system::Config for Runtime {
431 type Nonce = u64;
432 type AccountId = AccountId;
433 type Lookup = sp_runtime::traits::IdentityLookup<Self::AccountId>;
434 type Block = Block;
435 type AccountData = pallet_balances::AccountData<Balance>;
436}
437
438parameter_types! {
439 pub static ExistentialDeposit: Balance = 5;
440}
441
442#[derive_impl(pallet_balances::config_preludes::TestDefaultConfig)]
443impl pallet_balances::Config for Runtime {
444 type Balance = Balance;
445 type ExistentialDeposit = ExistentialDeposit;
446 type AccountStore = System;
447 type FreezeIdentifier = RuntimeFreezeReason;
448 type MaxFreezes = VariantCountOf<RuntimeFreezeReason>;
449 type RuntimeFreezeReason = RuntimeFreezeReason;
450}
451
452pub struct BalanceToU256;
453impl Convert<Balance, U256> for BalanceToU256 {
454 fn convert(n: Balance) -> U256 {
455 n.into()
456 }
457}
458
459pub struct U256ToBalance;
460impl Convert<U256, Balance> for U256ToBalance {
461 fn convert(n: U256) -> Balance {
462 n.try_into().unwrap()
463 }
464}
465
466pub struct RestrictMock;
467impl Contains<AccountId> for RestrictMock {
468 fn contains(who: &AccountId) -> bool {
469 RestrictedAccounts::get().contains(who)
470 }
471}
472
473parameter_types! {
474 pub static PostUnbondingPoolsWindow: u32 = 2;
475 pub static MaxMetadataLen: u32 = 2;
476 pub static CheckLevel: u8 = 255;
477 pub const PoolsPalletId: PalletId = PalletId(*b"py/nopls");
478}
479
480ord_parameter_types! {
481 pub const Admin: u128 = 42;
482}
483
484impl pools::Config for Runtime {
485 type RuntimeEvent = RuntimeEvent;
486 type WeightInfo = ();
487 type Currency = Balances;
488 type RuntimeFreezeReason = RuntimeFreezeReason;
489 type RewardCounter = RewardCounter;
490 type BalanceToU256 = BalanceToU256;
491 type U256ToBalance = U256ToBalance;
492 type StakeAdapter = adapter::DelegateStake<Self, StakingMock, DelegateMock>;
493 type PostUnbondingPoolsWindow = PostUnbondingPoolsWindow;
494 type PalletId = PoolsPalletId;
495 type MaxMetadataLen = MaxMetadataLen;
496 type MaxUnbonding = MaxUnbonding;
497 type MaxPointsToBalance = frame_support::traits::ConstU8<10>;
498 type AdminOrigin = EnsureSignedBy<Admin, AccountId>;
499 type BlockNumberProvider = System;
500 type Filter = RestrictMock;
501}
502
503type Block = frame_system::mocking::MockBlock<Runtime>;
504frame_support::construct_runtime!(
505 pub enum Runtime {
506 System: frame_system,
507 Balances: pallet_balances,
508 Pools: pools,
509 }
510);
511
512pub struct ExtBuilder {
513 members: Vec<(AccountId, Balance)>,
514 max_members: Option<u32>,
515 max_members_per_pool: Option<u32>,
516 global_max_commission: Option<Perbill>,
517}
518
519impl Default for ExtBuilder {
520 fn default() -> Self {
521 Self {
522 members: Default::default(),
523 max_members: Some(4),
524 max_members_per_pool: Some(3),
525 global_max_commission: Some(Perbill::from_percent(90)),
526 }
527 }
528}
529
530#[cfg_attr(feature = "fuzzing", allow(dead_code))]
531impl ExtBuilder {
532 pub fn add_members(mut self, members: Vec<(AccountId, Balance)>) -> Self {
534 self.members = members;
535 self
536 }
537
538 pub fn ed(self, ed: Balance) -> Self {
539 ExistentialDeposit::set(ed);
540 self
541 }
542
543 pub fn min_bond(self, min: Balance) -> Self {
544 StakingMinBond::set(min);
545 self
546 }
547
548 pub fn min_join_bond(self, min: Balance) -> Self {
549 MinJoinBondConfig::set(min);
550 self
551 }
552
553 pub fn with_check(self, level: u8) -> Self {
554 CheckLevel::set(level);
555 self
556 }
557
558 pub fn max_members(mut self, max: Option<u32>) -> Self {
559 self.max_members = max;
560 self
561 }
562
563 pub fn max_members_per_pool(mut self, max: Option<u32>) -> Self {
564 self.max_members_per_pool = max;
565 self
566 }
567
568 pub fn global_max_commission(mut self, commission: Option<Perbill>) -> Self {
569 self.global_max_commission = commission;
570 self
571 }
572
573 pub fn build(self) -> sp_io::TestExternalities {
574 sp_tracing::try_init_simple();
575 let mut storage =
576 frame_system::GenesisConfig::<Runtime>::default().build_storage().unwrap();
577
578 let _ = crate::GenesisConfig::<Runtime> {
579 min_join_bond: MinJoinBondConfig::get(),
580 min_create_bond: 2,
581 max_pools: Some(2),
582 max_members_per_pool: self.max_members_per_pool,
583 max_members: self.max_members,
584 global_max_commission: self.global_max_commission,
585 }
586 .assimilate_storage(&mut storage);
587
588 let mut ext = sp_io::TestExternalities::from(storage);
589
590 ext.execute_with(|| {
591 frame_system::Pallet::<Runtime>::set_block_number(1);
593
594 let amount_to_bond = Pools::depositor_min_bond();
596 Currency::set_balance(&10, amount_to_bond * 5);
597 assert_ok!(Pools::create(RawOrigin::Signed(10).into(), amount_to_bond, 900, 901, 902));
598 assert_ok!(Pools::set_metadata(RuntimeOrigin::signed(900), 1, vec![1, 1]));
599 let last_pool = LastPoolId::<Runtime>::get();
600 for (account_id, bonded) in self.members {
601 <Runtime as Config>::Currency::set_balance(&account_id, bonded * 2);
602 assert_ok!(Pools::join(RawOrigin::Signed(account_id).into(), bonded, last_pool));
603 }
604 });
605
606 ext
607 }
608
609 pub fn build_and_execute(self, test: impl FnOnce()) {
610 self.build().execute_with(|| {
611 test();
612 Pools::do_try_state(CheckLevel::get()).unwrap();
613 })
614 }
615}
616
617pub fn unsafe_set_state(pool_id: PoolId, state: PoolState) {
618 BondedPools::<Runtime>::try_mutate(pool_id, |maybe_bonded_pool| {
619 maybe_bonded_pool.as_mut().ok_or(()).map(|bonded_pool| {
620 bonded_pool.state = state;
621 })
622 })
623 .unwrap()
624}
625
626parameter_types! {
627 storage PoolsEvents: u32 = 0;
628 storage BalancesEvents: u32 = 0;
629}
630
631pub fn run_blocks(n: u64) {
633 let current_block = System::block_number();
634 System::run_to_block::<AllPalletsWithSystem>(n + current_block);
635}
636
637pub fn pool_events_since_last_call() -> Vec<super::Event<Runtime>> {
639 let events = System::events()
640 .into_iter()
641 .map(|r| r.event)
642 .filter_map(|e| if let RuntimeEvent::Pools(inner) = e { Some(inner) } else { None })
643 .collect::<Vec<_>>();
644 let already_seen = PoolsEvents::get();
645 PoolsEvents::set(&(events.len() as u32));
646 events.into_iter().skip(already_seen as usize).collect()
647}
648
649pub fn balances_events_since_last_call() -> Vec<pallet_balances::Event<Runtime>> {
651 let events = System::events()
652 .into_iter()
653 .map(|r| r.event)
654 .filter_map(|e| if let RuntimeEvent::Balances(inner) = e { Some(inner) } else { None })
655 .collect::<Vec<_>>();
656 let already_seen = BalancesEvents::get();
657 BalancesEvents::set(&(events.len() as u32));
658 events.into_iter().skip(already_seen as usize).collect()
659}
660
661pub fn fully_unbond_permissioned(member: AccountId) -> DispatchResult {
663 let points = PoolMembers::<Runtime>::get(member)
664 .map(|d| d.active_points())
665 .unwrap_or_default();
666 Pools::unbond(RuntimeOrigin::signed(member), member, points)
667}
668
669pub fn pending_rewards_for_delegator(delegator: AccountId) -> Balance {
670 let member = PoolMembers::<T>::get(delegator).unwrap();
671 let bonded_pool = BondedPools::<T>::get(member.pool_id).unwrap();
672 let reward_pool = RewardPools::<T>::get(member.pool_id).unwrap();
673
674 assert!(!bonded_pool.points.is_zero());
675
676 let commission = bonded_pool.commission.current();
677 let current_rc = reward_pool
678 .current_reward_counter(member.pool_id, bonded_pool.points, commission)
679 .unwrap()
680 .0;
681
682 member.pending_rewards(current_rc).unwrap_or_default()
683}
684
685#[derive(PartialEq, Debug)]
686pub enum RewardImbalance {
687 Surplus(Balance),
689 Deficit(Balance),
691}
692
693pub fn pool_pending_rewards(pool: PoolId) -> Result<BalanceOf<T>, sp_runtime::DispatchError> {
694 let bonded_pool = BondedPools::<T>::get(pool).ok_or(Error::<T>::PoolNotFound)?;
695 let reward_pool = RewardPools::<T>::get(pool).ok_or(Error::<T>::PoolNotFound)?;
696
697 let current_rc = if !bonded_pool.points.is_zero() {
698 let commission = bonded_pool.commission.current();
699 reward_pool.current_reward_counter(pool, bonded_pool.points, commission)?.0
700 } else {
701 Default::default()
702 };
703
704 Ok(PoolMembers::<T>::iter()
705 .filter(|(_, d)| d.pool_id == pool)
706 .map(|(_, d)| d.pending_rewards(current_rc).unwrap_or_default())
707 .fold(0u32.into(), |acc: BalanceOf<T>, x| acc.saturating_add(x)))
708}
709
710pub fn reward_imbalance(pool: PoolId) -> RewardImbalance {
711 let pending_rewards = pool_pending_rewards(pool).expect("pool should exist");
712 let current_balance = RewardPool::<Runtime>::current_balance(pool);
713
714 if pending_rewards > current_balance {
715 RewardImbalance::Deficit(pending_rewards - current_balance)
716 } else {
717 RewardImbalance::Surplus(current_balance - pending_rewards)
718 }
719}
720
721pub fn set_pool_balance(who: AccountId, amount: Balance) {
722 StakingMock::set_bonded_balance(who, amount);
723 DelegateMock::set_agent_balance(who, amount);
724}
725
726pub fn member_delegation(who: AccountId) -> Balance {
727 <T as Config>::StakeAdapter::member_delegation_balance(Member::from(who))
728 .expect("who must be a pool member")
729}
730
731pub fn pool_balance(id: PoolId) -> Balance {
732 <T as Config>::StakeAdapter::total_balance(Pool::from(Pools::generate_bonded_account(id)))
733 .expect("who must be a bonded pool account")
734}
735
736pub fn add_to_restrict_list(who: &AccountId) {
737 if !RestrictedAccounts::get().contains(who) {
738 RestrictedAccounts::mutate(|l| l.push(*who));
739 }
740}
741
742pub fn remove_from_restrict_list(who: &AccountId) {
743 RestrictedAccounts::mutate(|l| l.retain(|x| x != who));
744}
745
746#[cfg(test)]
747mod test {
748 use super::*;
749 #[test]
750 fn u256_to_balance_convert_works() {
751 assert_eq!(U256ToBalance::convert(0u32.into()), Zero::zero());
752 assert_eq!(U256ToBalance::convert(Balance::max_value().into()), Balance::max_value())
753 }
754
755 #[test]
756 #[should_panic]
757 fn u256_to_balance_convert_panics_correctly() {
758 U256ToBalance::convert(U256::from(Balance::max_value()).saturating_add(1u32.into()));
759 }
760
761 #[test]
762 fn balance_to_u256_convert_works() {
763 assert_eq!(BalanceToU256::convert(0u32.into()), U256::zero());
764 assert_eq!(BalanceToU256::convert(Balance::max_value()), Balance::max_value().into())
765 }
766}