Add withdraw functionality for metastaking rewards

farm-staking/farm-staking/src/custom_rewards.rs

@@ -43,6 +43,25 @@ pub trait CustomRewardsModule:
         self.reward_capacity().update(|r| *r += payment_amount);
     }
 
+    #[payable("*")]
+    #[endpoint(withdrawRewards)]
+    fn withdraw_rewards(&self, withdraw_amount: BigUint) {
+        self.require_caller_has_admin_permissions();
+
+        self.reward_capacity().update(|rewards| {
+            require!(
+                *rewards >= withdraw_amount,
+                "Not enough rewards to withdraw"
+            );
+
+            *rewards -= withdraw_amount.clone()
+        });
+
+        let caller = self.blockchain().get_caller();
+        let reward_token_id = self.reward_token_id().get();
+        self.send_tokens_non_zero(&caller, &reward_token_id, 0, &withdraw_amount);
+    }
+
     #[endpoint(endProduceRewards)]
     fn end_produce_rewards(&self) {
         self.require_caller_has_admin_permissions();

farm-staking/farm-staking/src/lib.rs

@@ -123,7 +123,7 @@ pub trait FarmStaking:
 
         (merged_farm_token, boosted_rewards_payment).into()
     }
-
+    
     #[view(calculateRewardsForGivenPosition)]
     fn calculate_rewards_for_given_position(
         &self,

farm-staking/farm-staking/tests/farm_staking_setup/mod.rs

@@ -416,6 +416,18 @@ where
             .assert_ok();
     }
 
+    pub fn check_rewards_capacity(&mut self, expected_farm_token_supply: u64) {
+        self.b_mock
+            .execute_query(&self.farm_wrapper, |sc| {
+                let actual_farm_supply = sc.reward_capacity().get();
+                assert_eq!(
+                    managed_biguint!(expected_farm_token_supply),
+                    actual_farm_supply
+                );
+            })
+            .assert_ok();
+    }
+
     pub fn allow_external_claim_rewards(&mut self, user: &Address) {
         self.b_mock
             .execute_tx(user, &self.farm_wrapper, &rust_biguint!(0), |sc| {
@@ -490,4 +502,17 @@ where
             )
             .assert_ok();
     }
+
+    pub fn withdraw_rewards(&mut self, withdraw_amount: &RustBigUint) {
+        self.b_mock
+            .execute_tx(
+                &self.owner_address,
+                &self.farm_wrapper,
+                &rust_biguint!(0),
+                |sc| {
+                    sc.withdraw_rewards(withdraw_amount.into());
+                },
+            )
+            .assert_ok();
+    }
 }

farm-staking/farm-staking/tests/farm_staking_test.rs

@@ -1,6 +1,8 @@
 #![allow(deprecated)]
 
-use multiversx_sc_scenario::{rust_biguint, whitebox_legacy::TxTokenTransfer, DebugApi};
+use multiversx_sc_scenario::{
+    rust_biguint, whitebox_legacy::TxTokenTransfer, DebugApi,
+};
 
 pub mod farm_staking_setup;
 use farm_staking::{
@@ -231,3 +233,19 @@ fn test_unbond() {
         USER_TOTAL_RIDE_TOKENS + expected_rewards,
     );
 }
+
+#[test]
+fn test_withdraw_rewards() {
+    DebugApi::dummy();
+    let mut farm_setup =
+        FarmStakingSetup::new(farm_staking::contract_obj, energy_factory::contract_obj);
+
+    let initial_rewards_capacity = 1_000_000_000_000u64;
+    farm_setup.check_rewards_capacity(initial_rewards_capacity);
+
+    let withdraw_amount = rust_biguint!(TOTAL_REWARDS_AMOUNT);
+    farm_setup.withdraw_rewards(&withdraw_amount);
+
+    let final_rewards_capacity = 0u64;
+    farm_setup.check_rewards_capacity(final_rewards_capacity);
+}