From 6d68ec7c3958d08474eeb05bb18379548eea44e0 Mon Sep 17 00:00:00 2001
From: drcpu <drcpu@protonmail.com>
Date: Wed, 29 Jan 2025 10:52:13 +0000
Subject: [PATCH] chore(timelock): fix environment-based timelock check for
 tests

---
 data_structures/src/transaction_factory.rs | 20 +++++++++++---------
 validations/src/validations.rs             |  9 +++++----
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/data_structures/src/transaction_factory.rs b/data_structures/src/transaction_factory.rs
index 88a3d62a0..4d405dd7c 100644
--- a/data_structures/src/transaction_factory.rs
+++ b/data_structures/src/transaction_factory.rs
@@ -680,16 +680,18 @@ pub fn transaction_inputs_sum(
         // Verify that commits are only accepted after the time lock expired
         let (epoch_timestamp, _) = epoch_constants.epoch_timestamp(epoch)?;
         let vt_time_lock = i64::try_from(vt_output.time_lock)?;
-        if (get_environment() == Environment::Development
-            || get_environment() == Environment::Testnet)
-            && epoch > 22_000
-            && vt_time_lock > epoch_timestamp
-        {
-            return Err(TransactionError::TimeLock {
-                expected: vt_time_lock,
-                current: epoch_timestamp,
+        if vt_time_lock > epoch_timestamp {
+            if get_environment() == Environment::Mainnet
+                || ((get_environment() == Environment::Development
+                    || get_environment() == Environment::Testnet)
+                    && epoch > 22_000)
+            {
+                return Err(TransactionError::TimeLock {
+                    expected: vt_time_lock,
+                    current: epoch_timestamp,
+                }
+                .into());
             }
-            .into());
         } else {
             if !seen_output_pointers.insert(input.output_pointer()) {
                 // If the set already contained this output pointer
diff --git a/validations/src/validations.rs b/validations/src/validations.rs
index ac25d5909..91481d3b4 100644
--- a/validations/src/validations.rs
+++ b/validations/src/validations.rs
@@ -1498,11 +1498,12 @@ pub fn validate_unstake_transaction<'a>(
 
             // TODO: modify this to enable delegated staking with multiple withdrawer addresses on a single validator
             // TODO: remove this check for mainnet release
-            if (get_environment() == Environment::Development
-                || get_environment() == Environment::Testnet)
-                && epoch > 22_000
+            let nonce = stake_entry.first().map(|stake| stake.value.nonce).unwrap();
+            if get_environment() == Environment::Mainnet
+                || ((get_environment() == Environment::Development
+                    || get_environment() == Environment::Testnet)
+                    && epoch > 22_000)
             {
-                let nonce = stake_entry.first().map(|stake| stake.value.nonce).unwrap();
                 if ut_tx.body.nonce != nonce {
                     return Err(TransactionError::UnstakeInvalidNonce {
                         used: ut_tx.body.nonce,