diff --git a/rewards/src/lib.rs b/rewards/src/lib.rs index 0ef94ade5..f395f9506 100644 --- a/rewards/src/lib.rs +++ b/rewards/src/lib.rs @@ -274,22 +274,24 @@ impl Pallet { return; } - PoolInfos::::mutate(pool, |pool_info| { - let total_shares = U256::from(pool_info.total_shares.to_owned().saturated_into::()); - pool_info.rewards.iter_mut().for_each( - |(reward_currency, (total_reward, total_withdrawn_reward))| { - Self::claim_one( - withdrawn_rewards, - *reward_currency, - share.to_owned(), - total_reward.to_owned(), - total_shares, - total_withdrawn_reward, - who, - pool, - ); - }, - ); + PoolInfos::::mutate_exists(pool, |maybe_pool_info| { + if let Some(pool_info) = maybe_pool_info { + let total_shares = U256::from(pool_info.total_shares.to_owned().saturated_into::()); + pool_info.rewards.iter_mut().for_each( + |(reward_currency, (total_reward, total_withdrawn_reward))| { + Self::claim_one( + withdrawn_rewards, + *reward_currency, + share.to_owned(), + total_reward.to_owned(), + total_shares, + total_withdrawn_reward, + who, + pool, + ); + }, + ); + } }); } }); diff --git a/rewards/src/tests.rs b/rewards/src/tests.rs index f0c05bd25..fb80a28bb 100644 --- a/rewards/src/tests.rs +++ b/rewards/src/tests.rs @@ -121,6 +121,32 @@ fn claim_rewards_should_not_create_empty_records() { SharesAndWithdrawnRewards::::contains_key(&DOT_POOL, &ALICE), false ); + + PoolInfos::::mutate(DOT_POOL, |pool_info| { + pool_info.rewards.insert(NATIVE_COIN, (10_000, 0)); + }); + RewardsModule::add_share(&ALICE, &DOT_POOL, 100); + assert_eq!( + RewardsModule::pool_infos(DOT_POOL), + PoolInfo { + total_shares: 100, + rewards: vec![(NATIVE_COIN, (10_000, 0))].into_iter().collect() + } + ); + assert_eq!( + RewardsModule::shares_and_withdrawn_rewards(DOT_POOL, ALICE), + (100, vec![(NATIVE_COIN, 0)].into_iter().collect()) + ); + + PoolInfos::::remove(DOT_POOL); + assert_eq!(PoolInfos::::contains_key(DOT_POOL), false); + + RewardsModule::claim_rewards(&ALICE, &DOT_POOL); + assert_eq!(PoolInfos::::contains_key(&DOT_POOL), false); + assert_eq!( + RewardsModule::shares_and_withdrawn_rewards(DOT_POOL, ALICE), + (100, vec![(NATIVE_COIN, 0)].into_iter().collect()) + ); }) }