From 0eddd3e8d796854703dc35e6c3a69d21265b491d Mon Sep 17 00:00:00 2001
From: Raul <raul.mf86@gmail.com>
Date: Wed, 9 Oct 2024 02:34:48 -0300
Subject: [PATCH] remove need for env function in test

---
 contracts/script/GenerateAlloc.s.sol | 18 ++++++++++++++++--
 contracts/test/utils/Test.sol        |  1 +
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/contracts/script/GenerateAlloc.s.sol b/contracts/script/GenerateAlloc.s.sol
index 06aa6769..95a1ebf1 100644
--- a/contracts/script/GenerateAlloc.s.sol
+++ b/contracts/script/GenerateAlloc.s.sol
@@ -35,8 +35,8 @@ contract GenerateAlloc is Script {
     // Upgrade admin controls upgradeability (by being Owner of each ProxyAdmin),
     // protocol admin is Owner of precompiles (admin/governance methods).
     // To disable upgradeability, we transfer ProxyAdmin ownership to a dead address
-    address internal upgradeAdmin = vm.envAddress("UPGRADE_ADMIN_ADDRESS");
-    address internal protocolAdmin = vm.envAddress("ADMIN_ADDRESS");
+    address internal upgradeAdmin;
+    address internal protocolAdmin;
     string internal dumpPath = getDumpPath();
     bool public saveState = true;
     uint256 public constant MAINNET_CHAIN_ID = 0; // TBD
@@ -47,6 +47,13 @@ contract GenerateAlloc is Script {
         saveState = false;
     }
 
+    /// @notice call from Test.sol only
+    function setAdminAddresses(address upgrade, address protocol) external {
+        require(block.chainid == 31337, "Only for local tests");
+        upgradeAdmin = upgrade;
+        protocolAdmin = protocol;
+    }
+
     /// @notice path where alloc file will be stored
     function getDumpPath() internal view returns (string memory) {
         if (block.chainid == 1513) {
@@ -62,7 +69,14 @@ contract GenerateAlloc is Script {
 
     /// @notice main script method
     function run() public {
+        if (upgradeAdmin == address(0)) {
+            upgradeAdmin = vm.envAddress("UPGRADE_ADMIN_ADDRESS");
+        }
         require(upgradeAdmin != address(0), "upgradeAdmin not set");
+
+        if (protocolAdmin == address(0)) {
+            protocolAdmin = vm.envAddress("ADMIN_ADDRESS");
+        }
         require(protocolAdmin != address(0), "protocolAdmin not set");
 
         vm.startPrank(deployer);
diff --git a/contracts/test/utils/Test.sol b/contracts/test/utils/Test.sol
index 97c830a4..28446b72 100644
--- a/contracts/test/utils/Test.sol
+++ b/contracts/test/utils/Test.sol
@@ -25,6 +25,7 @@ contract Test is ForgeTest {
     function setUp() virtual public {
         GenerateAlloc initializer = new GenerateAlloc();
         initializer.disableStateDump(); // Faster tests. Don't call to verify JSON output
+        initializer.setAdminAddresses(upgradeAdmin, admin);
         initializer.run();
         ipTokenStaking = IPTokenStaking(Predeploys.Staking);
         ipTokenSlashing = IPTokenSlashing(Predeploys.Slashing);