Add claimBoostedRewards endpoint in Fees Collector

multiversx · Dec 7, 2023 · 26607bd · 26607bd
1 parent fdfa965
commit 26607bd
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 12 deletions.
diff --git a/energy-integration/fees-collector/src/lib.rs b/energy-integration/fees-collector/src/lib.rs
@@ -48,13 +48,24 @@ pub trait FeesCollector:
     fn upgrade(&self) {}
 
     #[endpoint(claimRewards)]
-    fn claim_rewards(
+    fn claim_rewards_endpoint(
         &self,
         opt_original_caller: OptionalValue<ManagedAddress>,
     ) -> PaymentsVec<Self::Api> {
         require!(self.not_paused(), "Cannot claim while paused");
 
-        self.accumulate_additional_locked_tokens();
+        let caller = self.blockchain().get_caller();
+        let original_caller = self.get_orig_caller_from_opt(&caller, opt_original_caller);
+
+        self.claim_rewards(caller, original_caller)
+    }
+
+    #[endpoint(claimBoostedRewards)]
+    fn claim_boosted_rewards(
+        &self,
+        opt_original_caller: OptionalValue<ManagedAddress>,
+    ) -> PaymentsVec<Self::Api> {
+        require!(self.not_paused(), "Cannot claim while paused");
 
         let original_caller = match opt_original_caller {
             OptionalValue::Some(user) => {
@@ -67,6 +78,16 @@ pub trait FeesCollector:
             OptionalValue::None => self.blockchain().get_caller(),
         };
 
+        self.claim_rewards(original_caller.clone(), original_caller)
+    }
+
+    fn claim_rewards(
+        &self,
+        caller: ManagedAddress,
+        original_caller: ManagedAddress,
+    ) -> PaymentsVec<Self::Api> {
+        self.accumulate_additional_locked_tokens();
+
         let wrapper = FeesCollectorWrapper::new();
         let mut rewards = self.claim_multi(&wrapper, &original_caller);
         if rewards.is_empty() {
@@ -90,14 +111,14 @@ pub trait FeesCollector:
         }
 
         if !rewards.is_empty() {
-            self.send().direct_multi(&original_caller, &rewards);
+            self.send().direct_multi(&caller, &rewards);
         }
 
         if total_locked_token_rewards_amount > 0 {
             let locked_rewards = self.lock_virtual(
                 self.get_base_token_id(),
                 total_locked_token_rewards_amount,
-                original_caller.clone(),
+                caller,
                 original_caller,
             );
 

diff --git a/energy-integration/fees-collector/tests/fees_collector_test_setup/mod.rs b/energy-integration/fees-collector/tests/fees_collector_test_setup/mod.rs
@@ -225,14 +225,14 @@ where
     pub fn claim(&mut self, user: &Address) -> TxResult {
         self.b_mock
             .execute_tx(user, &self.fc_wrapper, &rust_biguint!(0), |sc| {
-                let _ = sc.claim_rewards(OptionalValue::None);
+                let _ = sc.claim_rewards_endpoint(OptionalValue::None);
             })
     }
 
     pub fn claim_for_user(&mut self, owner: &Address, broker: &Address) -> TxResult {
         self.b_mock
             .execute_tx(broker, &self.fc_wrapper, &rust_biguint!(0), |sc| {
-                let _ = sc.claim_rewards(OptionalValue::Some(managed_address!(owner)));
+                let _ = sc.claim_boosted_rewards(OptionalValue::Some(managed_address!(owner)));
             })
     }
 

diff --git a/energy-integration/fees-collector/wasm/src/lib.rs b/energy-integration/fees-collector/wasm/src/lib.rs
@@ -5,9 +5,9 @@
 ////////////////////////////////////////////////////
 
 // Init:                                 1
-// Endpoints:                           38
+// Endpoints:                           39
 // Async Callback (empty):               1
-// Total number of exported functions:  40
+// Total number of exported functions:  41
 
 #![no_std]
 
@@ -23,7 +23,8 @@ multiversx_sc_wasm_adapter::endpoints! {
     (
         init => init
         upgrade => upgrade
-        claimRewards => claim_rewards
+        claimRewards => claim_rewards_endpoint
+        claimBoostedRewards => claim_boosted_rewards
         addKnownContracts => add_known_contracts
         removeKnownContracts => remove_known_contracts
         addKnownTokens => add_known_tokens

diff --git a/energy-integration/governance-v2/tests/gov_test_setup/mod.rs b/energy-integration/governance-v2/tests/gov_test_setup/mod.rs
@@ -121,7 +121,7 @@ where
                 &fees_collector_wrapper,
                 &rust_biguint!(0),
                 |sc| {
-                    sc.claim_rewards(OptionalValue::None);
+                    sc.claim_rewards_endpoint(OptionalValue::None);
                 },
             )
             .assert_ok();
@@ -132,7 +132,7 @@ where
                 &fees_collector_wrapper,
                 &rust_biguint!(0),
                 |sc| {
-                    sc.claim_rewards(OptionalValue::None);
+                    sc.claim_rewards_endpoint(OptionalValue::None);
                 },
             )
             .assert_ok();
@@ -143,7 +143,7 @@ where
                 &fees_collector_wrapper,
                 &rust_biguint!(0),
                 |sc| {
-                    sc.claim_rewards(OptionalValue::None);
+                    sc.claim_rewards_endpoint(OptionalValue::None);
                 },
             )
             .assert_ok();