From cd6dce50bdbda814254c55d55e732c9a5ff604cd Mon Sep 17 00:00:00 2001
From: 1xstj <106580853+1xstj@users.noreply.github.com>
Date: Fri, 24 Nov 2023 12:27:11 +0000
Subject: [PATCH] add tests

---
 pallets/roles/src/impls.rs | 21 +++++++++-----
 pallets/roles/src/mock.rs  | 13 +++++++--
 pallets/roles/src/tests.rs | 57 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 81 insertions(+), 10 deletions(-)

diff --git a/pallets/roles/src/impls.rs b/pallets/roles/src/impls.rs
index 5ac9f175e..6c3e51b7b 100644
--- a/pallets/roles/src/impls.rs
+++ b/pallets/roles/src/impls.rs
@@ -18,10 +18,10 @@ use super::*;
 use frame_support::{
 	log,
 	pallet_prelude::DispatchResult,
-	traits::{Currency, ExistenceRequirement::KeepAlive, OneSessionHandler},
+	traits::{Currency, OneSessionHandler},
 };
-use sp_runtime::{Percent, Saturating};
-use sp_std::ops::Div;
+use sp_runtime::{traits::CheckedDiv, Percent, Saturating};
+
 use tangle_primitives::{jobs::JobKey, traits::roles::RolesHandler};
 
 /// Implements RolesHandler for the pallet.
@@ -208,10 +208,17 @@ impl<T: Config> Pallet<T> {
 		let reward_distribution = T::ValidatorRewardDistribution::get();
 
 		let dist = reward_distribution.get_reward_distribution();
-		let tss_validator_reward =
-			dist.0.mul_floor(total_rewards).div((tss_validators.len() as u32).into());
-		let zksaas_validators_reward =
-			dist.1.mul_floor(total_rewards).div((zksaas_validators.len() as u32).into());
+
+		let tss_validator_reward = dist
+			.0
+			.mul_floor(total_rewards)
+			.checked_div(&(tss_validators.len() as u32).into())
+			.unwrap_or(Zero::zero());
+		let zksaas_validators_reward = dist
+			.1
+			.mul_floor(total_rewards)
+			.checked_div(&(zksaas_validators.len() as u32).into())
+			.unwrap_or(Zero::zero());
 
 		log::debug!("Reward for tss validator : {:?}", tss_validator_reward);
 		log::debug!("Reward for zksaas validator : {:?}", zksaas_validators_reward);
diff --git a/pallets/roles/src/mock.rs b/pallets/roles/src/mock.rs
index 760badba1..b785a83da 100644
--- a/pallets/roles/src/mock.rs
+++ b/pallets/roles/src/mock.rs
@@ -27,10 +27,9 @@ use sp_core::H256;
 use sp_runtime::{
 	testing::{Header, UintAuthorityId},
 	traits::IdentityLookup,
-	BuildStorage, DispatchResult, Perbill,
+	BuildStorage, DispatchResult, Perbill, Percent,
 };
-use tangle_primitives::{jobs::*, traits::jobs::MPCHandler};
-
+use tangle_primitives::{jobs::*, roles::ValidatorRewardDistribution, traits::jobs::MPCHandler};
 pub type AccountId = u64;
 pub type Balance = u128;
 pub type BlockNumber = u64;
@@ -243,11 +242,19 @@ impl JobsHandler<AccountId> for MockJobsHandler {
 	}
 }
 
+parameter_types! {
+	pub InflationRewardPerEra: Balance = 10_000;
+	pub Reward : ValidatorRewardDistribution = ValidatorRewardDistribution::try_new(Percent::from_rational(1_u32,2_u32), Percent::from_rational(1_u32,2_u32)).unwrap();
+}
+
 impl Config for Runtime {
 	type RuntimeEvent = RuntimeEvent;
 	type JobsHandler = MockJobsHandler;
 	type MaxRolesPerAccount = ConstU32<2>;
 	type MPCHandler = MockMPCHandler;
+	type InflationRewardPerEra = InflationRewardPerEra;
+	type AuthorityId = UintAuthorityId;
+	type ValidatorRewardDistribution = Reward;
 	type WeightInfo = ();
 }
 
diff --git a/pallets/roles/src/tests.rs b/pallets/roles/src/tests.rs
index 8dbfe0b2d..302c88305 100644
--- a/pallets/roles/src/tests.rs
+++ b/pallets/roles/src/tests.rs
@@ -231,3 +231,60 @@ fn test_unbound_funds_should_work_if_no_role_assigned() {
 		})]);
 	});
 }
+
+#[test]
+fn test_reward_dist_works_as_expected_with_one_validator() {
+	new_test_ext_raw_authorities(vec![1, 2, 3, 4]).execute_with(|| {
+		assert_eq!(Balances::free_balance(1), 20_000);
+
+		// Roles user is interested in re-staking.
+		let role_records = vec![
+			RoleStakingRecord {
+				metadata: RoleTypeMetadata::Tss(Default::default()),
+				re_staked: 2500,
+			},
+			RoleStakingRecord {
+				metadata: RoleTypeMetadata::ZkSaas(Default::default()),
+				re_staked: 2500,
+			},
+		];
+
+		assert_ok!(Roles::assign_roles(RuntimeOrigin::signed(1), role_records));
+
+		// The reward is 100, we have 5 authorities
+		assert_ok!(Roles::distribute_rewards());
+
+		// ensure the distribution is correct
+		assert_eq!(Balances::free_balance(1), 20_000 + 10_000);
+	});
+}
+
+#[test]
+fn test_reward_dist_works_as_expected_with_multiple_validator() {
+	new_test_ext_raw_authorities(vec![1, 2, 3, 4]).execute_with(|| {
+		let _reward_amount = 10_000;
+		assert_eq!(Balances::free_balance(1), 20_000);
+
+		// Roles user is interested in re-staking.
+		let role_records = vec![
+			RoleStakingRecord {
+				metadata: RoleTypeMetadata::Tss(Default::default()),
+				re_staked: 2500,
+			},
+			RoleStakingRecord {
+				metadata: RoleTypeMetadata::ZkSaas(Default::default()),
+				re_staked: 2500,
+			},
+		];
+
+		assert_ok!(Roles::assign_roles(RuntimeOrigin::signed(1), role_records.clone()));
+		assert_ok!(Roles::assign_roles(RuntimeOrigin::signed(2), role_records));
+
+		// The reward is 100, we have 5 authorities
+		assert_ok!(Roles::distribute_rewards());
+
+		// ensure the distribution is correct
+		assert_eq!(Balances::free_balance(1), 20_000 + 5000);
+		assert_eq!(Balances::free_balance(2), 20_000 + 5000);
+	});
+}