Add set_vesting call to Claims. Add unit tests

xx-labs · Dec 29, 2021 · 12b14db · 12b14db
1 parent 8e4a0cb
commit 12b14db
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 0 deletions.
diff --git a/claims/src/lib.rs b/claims/src/lib.rs
@@ -413,6 +413,23 @@ decl_module! {
 			));
 			Ok(Pays::No.into())
 		}
+
+        #[weight = (
+            T::WeightInfo::set_vesting(),
+            DispatchClass::Operational,
+            Pays::No
+        )]
+        fn set_vesting(origin,
+            who: EthereumAddress,
+            vesting_schedules: Option<Vec<(BalanceOf<T>, BalanceOf<T>, T::BlockNumber)>>,
+        ) {
+            T::MoveClaimOrigin::try_origin(origin).map(|_| ()).or_else(ensure_root)?;
+            if let Some(vs) = vesting_schedules {
+                <Vesting<T>>::insert(who, vs);
+            } else {
+                <Vesting<T>>::take(who);
+            }
+        }
 	}
 }
 
@@ -1243,6 +1260,30 @@ mod tests {
             );
         });
     }
+
+    #[test]
+    fn set_vesting_works() {
+        new_test_ext().execute_with(|| {
+            // Add vesting to a claim that doesn't have one
+            assert_eq!(Claims::vesting(&eth(&eve())), None);
+            assert_ok!(Claims::set_vesting(Origin::signed(6), eth(&eve()), Some(vec![(300, 10, 0)])));
+            assert_eq!(Claims::vesting(&eth(&eve())), Some(vec![(300, 10, 0)]));
+            // Check that nothing bad happens when remove vesting from a claim that doesn't have one
+            assert_eq!(Claims::vesting(&eth(&bob())), None);
+            assert_ok!(Claims::set_vesting(Origin::signed(6), eth(&bob()), None));
+            assert_eq!(Claims::vesting(&eth(&bob())), None);
+            // Remove existing vesting from a claim
+            assert_eq!(Claims::vesting(&eth(&alice())), Some(vec![(50, 10, 1)]));
+            assert_ok!(Claims::set_vesting(Origin::signed(6), eth(&alice()), None));
+            assert_eq!(Claims::vesting(&eth(&alice())), None);
+            // Modify existing vesting
+            assert_eq!(Claims::vesting(&eth(&charlie())), Some(vec![(50, 10, 0), (450, 10, 0)]));
+            assert_ok!(Claims::set_vesting(Origin::signed(6), eth(&charlie()), Some(vec![(100, 10, 0), (400, 10, 0)])));
+            assert_eq!(Claims::vesting(&eth(&charlie())), Some(vec![(100, 10, 0), (400, 10, 0)]));
+            // Can't call set vesting from non privileged origin
+            assert_noop!(Claims::set_vesting(Origin::signed(42), eth(&alice()), None), BadOrigin);
+        });
+    }
 }
 
 #[cfg(feature = "runtime-benchmarks")]

diff --git a/claims/src/weights.rs b/claims/src/weights.rs
@@ -25,6 +25,7 @@ pub trait WeightInfo {
     fn claim_attest() -> Weight;
     fn attest() -> Weight;
     fn move_claim() -> Weight;
+    fn set_vesting() -> Weight;
 }
 
 pub struct TestWeightInfo;
@@ -34,6 +35,7 @@ impl WeightInfo for TestWeightInfo {
     fn claim_attest() -> Weight { 0 }
     fn attest() -> Weight { 0 }
     fn move_claim() -> Weight { 0 }
+    fn set_vesting() -> Weight { 0 }
 }
 
 /// Weight functions for claims.
@@ -64,4 +66,8 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
             .saturating_add(T::DbWeight::get().reads(4 as Weight))
             .saturating_add(T::DbWeight::get().writes(7 as Weight))
     }
+    fn set_vesting() -> Weight {
+        (24_376_000 as Weight)
+            .saturating_add(T::DbWeight::get().writes(1 as Weight))
+    }
 }
diff --git a/runtime/xxnetwork/src/weights/claims.rs b/runtime/xxnetwork/src/weights/claims.rs
@@ -87,4 +87,9 @@ impl<T: frame_system::Config> claims::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(4 as Weight))
 			.saturating_add(T::DbWeight::get().writes(7 as Weight))
 	}
+	// Storage: Claims Vesting (r:0 w:1)
+	fn set_vesting() -> Weight {
+		(24_376_000 as Weight)
+			.saturating_add(T::DbWeight::get().writes(1 as Weight))
+	}
 }