diff --git a/farm-staking/farm-staking/src/custom_rewards.rs b/farm-staking/farm-staking/src/custom_rewards.rs index de7fd29a0..8e4d538dc 100644 --- a/farm-staking/farm-staking/src/custom_rewards.rs +++ b/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(); diff --git a/farm-staking/farm-staking/src/lib.rs b/farm-staking/farm-staking/src/lib.rs index 33f5bbd1a..ac1f23727 100644 --- a/farm-staking/farm-staking/src/lib.rs +++ b/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, diff --git a/farm-staking/farm-staking/tests/farm_staking_setup/mod.rs b/farm-staking/farm-staking/tests/farm_staking_setup/mod.rs index 51841149f..2a1109add 100644 --- a/farm-staking/farm-staking/tests/farm_staking_setup/mod.rs +++ b/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(); + } } diff --git a/farm-staking/farm-staking/tests/farm_staking_test.rs b/farm-staking/farm-staking/tests/farm_staking_test.rs index ed0925510..6de87b547 100644 --- a/farm-staking/farm-staking/tests/farm_staking_test.rs +++ b/farm-staking/farm-staking/tests/farm_staking_test.rs @@ -1,10 +1,12 @@ #![allow(deprecated)] -use multiversx_sc_scenario::{rust_biguint, whitebox_legacy::TxTokenTransfer, DebugApi}; +use multiversx_sc_scenario::{ + managed_address, rust_biguint, whitebox_legacy::TxTokenTransfer, DebugApi, +}; pub mod farm_staking_setup; use farm_staking::{ - custom_rewards::{BLOCKS_IN_YEAR, MAX_PERCENT}, + custom_rewards::{CustomRewardsModule, BLOCKS_IN_YEAR, MAX_PERCENT}, token_attributes::UnbondSftAttributes, }; use farm_staking_setup::*; @@ -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); +}