From 067064bc48f6cc6081b7c09bf73bf1263924c106 Mon Sep 17 00:00:00 2001
From: Sorin Petreasca <petreasca.sorin@gmail.com>
Date: Wed, 8 May 2024 00:45:01 +0300
Subject: [PATCH 1/2] Energy factory unit test audit update

---
 .../tests/simple_lock_energy_test.rs          | 80 ++++++++++++++++++-
 1 file changed, 79 insertions(+), 1 deletion(-)

diff --git a/locked-asset/energy-factory/tests/simple_lock_energy_test.rs b/locked-asset/energy-factory/tests/simple_lock_energy_test.rs
index 6d23bd1ef..8c0655476 100644
--- a/locked-asset/energy-factory/tests/simple_lock_energy_test.rs
+++ b/locked-asset/energy-factory/tests/simple_lock_energy_test.rs
@@ -2,7 +2,9 @@
 
 mod energy_factory_setup;
 
-use energy_factory::energy::EnergyModule;
+use energy_factory::{
+    energy::EnergyModule, locked_token_transfer::LockedTokenTransferModule, SimpleLockEnergy,
+};
 use energy_factory_setup::*;
 use multiversx_sc::types::BigUint;
 use simple_lock::locked_token::LockedTokenAttributes;
@@ -489,3 +491,79 @@ fn energy_deplete_test() {
         })
         .assert_ok();
 }
+
+#[test]
+fn extend_lock_period_endpoint_test() {
+    let mut setup = SimpleLockEnergySetup::new(energy_factory::contract_obj);
+    let first_user = setup.first_user.clone();
+    let random_user = setup.b_mock.create_user_account(&rust_biguint!(0u64));
+
+    let current_epoch = 1;
+    setup.b_mock.set_block_epoch(current_epoch);
+
+    setup
+        .lock(
+            &first_user,
+            BASE_ASSET_TOKEN_ID,
+            USER_BALANCE,
+            LOCK_OPTIONS[0],
+        )
+        .assert_ok();
+
+    setup
+        .b_mock
+        .execute_esdt_transfer(
+            &first_user,
+            &setup.sc_wrapper,
+            LOCKED_TOKEN_ID,
+            1,
+            &rust_biguint!(USER_BALANCE),
+            |sc| {
+                sc.extend_lock_period(100, managed_address!(&first_user));
+            },
+        )
+        .assert_user_error("Invalid lock choice");
+
+    setup
+        .b_mock
+        .execute_esdt_transfer(
+            &first_user,
+            &setup.sc_wrapper,
+            LOCKED_TOKEN_ID,
+            1,
+            &rust_biguint!(USER_BALANCE),
+            |sc| {
+                sc.extend_lock_period(LOCK_OPTIONS[0], managed_address!(&first_user));
+            },
+        )
+        .assert_user_error("May not call this endpoint. Use lockTokens instead");
+
+    let energy_per_epoch = rust_biguint!(USER_BALANCE);
+    let energy_before = setup.get_user_energy(&first_user);
+    assert_eq!(
+        energy_before,
+        LOCK_OPTIONS[0] * energy_per_epoch.clone() - energy_per_epoch.clone()
+    );
+
+    setup
+        .b_mock
+        .execute_esdt_transfer(
+            &first_user,
+            &setup.sc_wrapper,
+            LOCKED_TOKEN_ID,
+            1,
+            &rust_biguint!(USER_BALANCE),
+            |sc| {
+                sc.token_transfer_whitelist()
+                    .add(&managed_address!(&first_user));
+                sc.extend_lock_period(LOCK_OPTIONS[1], managed_address!(&first_user));
+            },
+        )
+        .assert_ok();
+
+    let energy_after = setup.get_user_energy(&first_user);
+    assert_eq!(
+        energy_after,
+        LOCK_OPTIONS[1] * energy_per_epoch.clone() - energy_per_epoch.clone()
+    );
+}

From 867a88d82192d515d9a3c01b6324073d700febec Mon Sep 17 00:00:00 2001
From: Sorin Petreasca <petreasca.sorin@gmail.com>
Date: Thu, 9 May 2024 17:15:59 +0300
Subject: [PATCH 2/2] clippy fix

---
 locked-asset/energy-factory/tests/simple_lock_energy_test.rs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/locked-asset/energy-factory/tests/simple_lock_energy_test.rs b/locked-asset/energy-factory/tests/simple_lock_energy_test.rs
index 8c0655476..202630958 100644
--- a/locked-asset/energy-factory/tests/simple_lock_energy_test.rs
+++ b/locked-asset/energy-factory/tests/simple_lock_energy_test.rs
@@ -496,7 +496,6 @@ fn energy_deplete_test() {
 fn extend_lock_period_endpoint_test() {
     let mut setup = SimpleLockEnergySetup::new(energy_factory::contract_obj);
     let first_user = setup.first_user.clone();
-    let random_user = setup.b_mock.create_user_account(&rust_biguint!(0u64));
 
     let current_epoch = 1;
     setup.b_mock.set_block_epoch(current_epoch);