adjustUserEnergy functionality

locked-asset/energy-factory/src/lib.rs

@@ -230,4 +230,33 @@ pub trait SimpleLockEnergy:
 
         output_tokens
     }
+
+    #[only_owner]
+    #[endpoint(adjustUserEnergy)]
+    fn adjust_user_energy(
+        &self,
+        args: MultiValueEncoded<MultiValue3<ManagedAddress, BigInt, BigInt>>,
+    ) {
+        for arg in args {
+            let (user, energy_amount, token_amount) = arg.into_tuple();
+            require!(!self.user_energy(&user).is_empty(), "User energy not found");
+            let old_energy = self.get_updated_energy_entry_for_user(&user);
+            let new_energy_amount = old_energy.get_energy_amount_raw() + &energy_amount;
+            let new_total_locked_tokens = if token_amount >= 0 {
+                old_energy.get_total_locked_tokens() + &token_amount.magnitude()
+            } else {
+                let token_amount_magnitude = token_amount.magnitude();
+                require!(
+                    old_energy.get_total_locked_tokens() >= &token_amount_magnitude,
+                    "Insufficient locked tokens"
+                );
+                old_energy.get_total_locked_tokens() - &token_amount_magnitude
+            };
+
+            let current_epoch = self.blockchain().get_block_epoch();
+            let new_energy = Energy::new(new_energy_amount, current_epoch, new_total_locked_tokens);
+
+            self.set_energy_entry(&user, new_energy);
+        }
+    }
 }

locked-asset/energy-factory/tests/energy_factory_setup/mod.rs

@@ -10,7 +10,7 @@ use energy_factory::{
 use multiversx_sc::{
     codec::multi_types::OptionalValue,
     storage::mappers::StorageTokenWrapper,
-    types::{Address, EsdtLocalRole, MultiValueEncoded},
+    types::{Address, BigInt, EsdtLocalRole, MultiValueEncoded},
 };
 use multiversx_sc_modules::pause::PauseModule;
 use multiversx_sc_scenario::{
@@ -19,6 +19,8 @@ use multiversx_sc_scenario::{
     whitebox_legacy::{BlockchainStateWrapper, ContractObjWrapper},
     DebugApi,
 };
+
+use num_bigint::Sign;
 use simple_lock::locked_token::LockedTokenModule;
 
 use unbond_sc_mock::*;
@@ -272,6 +274,47 @@ where
 
         result
     }
+
+    pub fn get_user_energy_raw(&mut self, user: &Address) -> num_bigint::BigInt {
+        let mut result = num_bigint::BigInt::from_biguint(Sign::NoSign, rust_biguint!(0));
+        self.b_mock
+            .execute_query(&self.sc_wrapper, |sc| {
+                let user_energy = sc.get_updated_energy_entry_for_user(&managed_address!(user));
+                result = to_rust_bigint(user_energy.get_energy_amount_raw().clone());
+            })
+            .assert_ok();
+
+        result
+    }
+
+    pub fn get_user_locked_tokens(&mut self, user: &Address) -> num_bigint::BigUint {
+        let mut result = rust_biguint!(0);
+        self.b_mock
+            .execute_query(&self.sc_wrapper, |sc| {
+                let user_energy = sc.get_updated_energy_entry_for_user(&managed_address!(user));
+                result = to_rust_biguint(user_energy.get_total_locked_tokens().clone());
+            })
+            .assert_ok();
+
+        result
+    }
+
+    pub fn adjust_user_energy(&mut self, user: &Address, energy_amount: i64, token_amount: i64) {
+        self.b_mock
+            .execute_tx(&self.owner, &self.sc_wrapper, &rust_biguint!(0), |sc| {
+                let mut args = MultiValueEncoded::new();
+                args.push(
+                    (
+                        managed_address!(user),
+                        BigInt::from(energy_amount),
+                        BigInt::from(token_amount),
+                    )
+                        .into(),
+                );
+                sc.adjust_user_energy(args);
+            })
+            .assert_ok();
+    }
 }
 
 pub fn to_rust_biguint(
@@ -280,6 +323,16 @@ pub fn to_rust_biguint(
     num_bigint::BigUint::from_bytes_be(managed_biguint.to_bytes_be().as_slice())
 }
 
+pub fn to_rust_bigint(
+    managed_biguint: multiversx_sc::types::BigInt<DebugApi>,
+) -> num_bigint::BigInt {
+    if managed_biguint < 0 {
+        num_bigint::BigInt::from_biguint(Sign::Minus, to_rust_biguint(managed_biguint.magnitude()))
+    } else {
+        num_bigint::BigInt::from_biguint(Sign::Plus, to_rust_biguint(managed_biguint.magnitude()))
+    }
+}
+
 pub fn to_start_of_month(unlock_epoch: u64) -> u64 {
     unlock_epoch - unlock_epoch % 30
 }

locked-asset/energy-factory/tests/simple_lock_energy_test.rs

@@ -7,6 +7,8 @@ use energy_factory::{
 };
 use energy_factory_setup::*;
 use multiversx_sc::types::BigUint;
+use num_bigint::BigInt;
+use num_traits::{FromPrimitive, Zero};
 use simple_lock::locked_token::LockedTokenAttributes;
 
 use multiversx_sc_scenario::{
@@ -566,3 +568,163 @@ fn extend_lock_period_endpoint_test() {
         LOCK_OPTIONS[1] * energy_per_epoch.clone() - energy_per_epoch.clone()
     );
 }
+
+#[test]
+fn adjust_user_energy_test() {
+    let mut setup = SimpleLockEnergySetup::new(energy_factory::contract_obj);
+    let user = setup.first_user.clone();
+    let user_balance = 1_000_000u64;
+
+    // Initial lock to create user energy
+    setup
+        .lock(&user, BASE_ASSET_TOKEN_ID, user_balance, LOCK_OPTIONS[0])
+        .assert_ok();
+
+    // Check initial energy
+    let initial_locked_tokens = rust_biguint!(user_balance);
+    let initial_energy = initial_locked_tokens.clone() * LOCK_OPTIONS[0];
+    assert_eq!(setup.get_user_locked_tokens(&user), initial_locked_tokens);
+    assert_eq!(setup.get_user_energy(&user), initial_energy);
+
+    // Case 1: Positive energy, positive token amount
+    let adjustment1_energy = 1_000_000i64;
+    let adjustment1_tokens = 500_000i64;
+    setup.adjust_user_energy(&user, adjustment1_energy, adjustment1_tokens);
+
+    let locked_tokens1 = setup.get_user_locked_tokens(&user);
+    let energy1 = setup.get_user_energy(&user);
+    assert_eq!(
+        locked_tokens1,
+        &initial_locked_tokens + &rust_biguint!(adjustment1_tokens)
+    );
+    assert_eq!(
+        energy1,
+        &initial_energy + &rust_biguint!(adjustment1_energy)
+    );
+
+    // Case 2: Positive energy, negative token amount
+    let adjustment2_energy = 800_000i64;
+    let adjustment2_tokens = -300_000i64;
+    setup.adjust_user_energy(&user, adjustment2_energy, adjustment2_tokens);
+
+    let locked_tokens2 = setup.get_user_locked_tokens(&user);
+    let energy2 = setup.get_user_energy(&user);
+    assert_eq!(
+        locked_tokens2,
+        &locked_tokens1 - &rust_biguint!(adjustment2_tokens.unsigned_abs())
+    );
+    assert_eq!(energy2, &energy1 + &rust_biguint!(adjustment2_energy));
+
+    // Case 3: Negative energy, positive token amount
+    let adjustment3_energy = -1_500_000i64;
+    let adjustment3_tokens = 200_000i64;
+    setup.adjust_user_energy(&user, adjustment3_energy, adjustment3_tokens);
+
+    let locked_tokens3 = setup.get_user_locked_tokens(&user);
+    let energy3 = setup.get_user_energy(&user);
+    assert_eq!(
+        locked_tokens3,
+        &locked_tokens2 + &rust_biguint!(adjustment3_tokens)
+    );
+    assert_eq!(
+        energy3,
+        &energy2 - &rust_biguint!(adjustment3_energy.unsigned_abs())
+    );
+
+    // Case 4: Negative energy, negative token amount
+    let adjustment4_energy = -700_000i64;
+    let adjustment4_tokens = -100_000i64;
+    setup.adjust_user_energy(&user, adjustment4_energy, adjustment4_tokens);
+
+    let locked_tokens4 = setup.get_user_locked_tokens(&user);
+    let energy4 = setup.get_user_energy(&user);
+    assert_eq!(
+        locked_tokens4,
+        &locked_tokens3 - &rust_biguint!(adjustment4_tokens.unsigned_abs())
+    );
+    assert_eq!(
+        energy4,
+        &energy3 - &rust_biguint!(adjustment4_energy.unsigned_abs())
+    );
+
+    // Bring energy to negative value
+    let adjustment5_energy = -500_000_000i64; // Adjust to negative energy
+    let adjustment5_tokens = 0i64; // Positive token amount
+    setup.adjust_user_energy(&user, adjustment5_energy, adjustment5_tokens);
+
+    let locked_tokens5 = setup.get_user_locked_tokens(&user);
+    let energy5 = setup.get_user_energy_raw(&user);
+    assert_eq!(
+        locked_tokens5,
+        &locked_tokens4 + &rust_biguint!(adjustment5_tokens)
+    );
+    assert!(energy5 < BigInt::zero()); // Energy should be negative
+
+    // Replay Case 1 on negative energy
+    let adjustment6_energy = 500_000i64;
+    let adjustment6_tokens = 200_000i64;
+    setup.adjust_user_energy(&user, adjustment6_energy, adjustment6_tokens);
+
+    let locked_tokens6 = setup.get_user_locked_tokens(&user);
+    let energy6 = setup.get_user_energy_raw(&user);
+    assert_eq!(
+        locked_tokens6,
+        &locked_tokens5 + &rust_biguint!(adjustment6_tokens)
+    );
+    assert_eq!(
+        energy6,
+        &energy5 + &BigInt::from_i64(adjustment6_energy).unwrap()
+    );
+    assert!(energy6 < BigInt::zero()); // Energy should remain negative
+
+    // Replay Case 2 on negative energy
+    let adjustment7_energy = 300_000i64;
+    let adjustment7_tokens = -100_000i64;
+    setup.adjust_user_energy(&user, adjustment7_energy, adjustment7_tokens);
+
+    let locked_tokens7 = setup.get_user_locked_tokens(&user);
+    let energy7 = setup.get_user_energy_raw(&user);
+    assert_eq!(
+        locked_tokens7,
+        &locked_tokens6 - &rust_biguint!(adjustment7_tokens.unsigned_abs())
+    );
+    assert_eq!(
+        energy7,
+        &energy6 + &BigInt::from_i64(adjustment7_energy).unwrap()
+    );
+    assert!(energy7 < BigInt::zero()); // Energy should remain negative
+
+    // Replay Case 3 on negative energy
+    let adjustment8_energy = -500_000i64;
+    let adjustment8_tokens = 100_000i64;
+    setup.adjust_user_energy(&user, adjustment8_energy, adjustment8_tokens);
+
+    let locked_tokens8 = setup.get_user_locked_tokens(&user);
+    let energy8 = setup.get_user_energy_raw(&user);
+    assert_eq!(
+        locked_tokens8,
+        &locked_tokens7 + &rust_biguint!(adjustment8_tokens)
+    );
+    assert_eq!(
+        energy8,
+        &energy7 + &BigInt::from_i64(adjustment8_energy).unwrap()
+    );
+    assert!(energy8 < BigInt::zero()); // Energy should remain negative
+
+    // Replay Case 4 on negative energy
+    let adjustment9_energy = -300_000i64;
+    let adjustment9_tokens = -50_000i64;
+    setup.adjust_user_energy(&user, adjustment9_energy, adjustment9_tokens);
+
+    let locked_tokens9 = setup.get_user_locked_tokens(&user);
+    let energy9 = setup.get_user_energy_raw(&user);
+    assert_eq!(
+        locked_tokens9,
+        &locked_tokens8 - &rust_biguint!(adjustment9_tokens.unsigned_abs())
+    );
+    assert_eq!(
+        energy9,
+        &energy8 + &BigInt::from_i64(adjustment9_energy).unwrap()
+    );
+    assert!(energy9 < BigInt::zero()); // Energy should remain negative
+}

locked-asset/energy-factory/wasm/src/lib.rs

@@ -6,9 +6,9 @@
 
 // Init:                                 1
 // Upgrade:                              1
-// Endpoints:                           33
+// Endpoints:                           34
 // Async Callback:                       1
-// Total number of exported functions:  36
+// Total number of exported functions:  37
 
 #![no_std]
 
@@ -23,6 +23,7 @@ multiversx_sc_wasm_adapter::endpoints! {
         lockTokens => lock_tokens_endpoint
         unlockTokens => unlock_tokens_endpoint
         extendLockPeriod => extend_lock_period
+        adjustUserEnergy => adjust_user_energy
         issueLockedToken => issue_locked_token
         getLockedTokenId => locked_token
         getBaseAssetTokenId => base_asset_token_id