enable M3 upgrade test for mainnet as well

yieldnest · Sep 29, 2024 · 2003475 · 2003475
1 parent ab5779a
commit 2003475
Show file tree

Hide file tree

Showing 3 changed files with 136 additions and 28 deletions.
diff --git a/test/integration/M3/Base.t.sol b/test/integration/M3/Base.t.sol
@@ -28,6 +28,8 @@ import {StakingNode} from "../../../src/StakingNode.sol";
 import {WithdrawalQueueManager} from "../../../src/WithdrawalQueueManager.sol";
 import {ynETHRedemptionAssetsVault} from "../../../src/ynETHRedemptionAssetsVault.sol";
 import {IStakingNode} from "../../../src/interfaces/IStakingNodesManager.sol";
+import {PlaceholderStakingNodesManager} from "./PlaceholderStakingNodesManager.sol";
+
 
 
 import "forge-std/console.sol";
@@ -111,7 +113,7 @@ contract Base is Test, Utils {
     }
 
     function upgradeYnToM3() internal {
-        if (block.chainid != 17000) return;
+        // if (block.chainid != 17000) return;
 
 
         uint256 totalAssets = yneth.totalAssets();
@@ -129,12 +131,49 @@ contract Base is Test, Utils {
 
         // upgrade stakingNodesManager
         {
+
+            address stakinNodesManagerImplementation;
+
+            if (block.chainid == 1) {
+
+                /*
+                 ██████╗  █████╗ ███╗   ██╗ ██████╗ ███████╗██████╗ 
+                 ██╔══██╗██╔══██╗████╗  ██║██╔════╝ ██╔════╝██╔══██╗
+                 ██║  ██║███████║██╔██╗ ██║██║  ███╗█████╗  ██████╔╝
+                 ██║  ██║██╔══██║██║╚██╗██║██║   ██║██╔══╝  ██╔══██╗
+                 ██████╔╝██║  ██║██║ ╚████║╚██████╔╝███████╗██║  ██║
+                 ╚═════╝ ╚═╝  ╚═╝╚═╝  ╚═══╝ ╚═════╝ ╚══════╝╚═╝  ╚═╝
+                */
+
+                // WARNING: This code is for testing purposes only and MUST be removed before deploying to mainnet.
+                // It uses a placeholder implementation that doesn't reflect the actual mainnet behavior.
+                // Keeping this in production could lead to severe security vulnerabilities and incorrect contract behavior.
+                // This logic auto-adjust for unverifiedStakedETH based on the the difference between podOwnerShares and pre-deposit balance.
+
+                // Compute deltas array
+                uint256[] memory deltas = new uint256[](stakingNodesManager.nodesLength());
+                for (uint256 i = 0; i < stakingNodesManager.nodesLength(); i++) {
+                    IStakingNode stakingNode = stakingNodesManager.nodes(i);
+                    uint256 podShares = uint256(IEigenPodManager(chainAddresses.eigenlayer.EIGENPOD_MANAGER_ADDRESS).podOwnerShares(address(stakingNode)));
+                    deltas[i] = preUpgradeState.stakingNodeBalances[i] - podShares;
+                    // Revert if delta is bigger than 32 ether
+                    require(deltas[i] <= 32 ether, "Delta exceeds 32 ether limit");
+                }
+
+                // Deploy PlaceholderStakingNodesManager
+                PlaceholderStakingNodesManager placeholderStakingNodesManager = new PlaceholderStakingNodesManager(deltas);
+                stakinNodesManagerImplementation = address(placeholderStakingNodesManager);
+            } else {
+                stakinNodesManagerImplementation = address(new StakingNodesManager());
+            }
+
+
             vm.startPrank(actors.admin.PROXY_ADMIN_OWNER);
             ProxyAdmin(
                 getTransparentUpgradeableProxyAdminAddress(address(stakingNodesManager))
             ).upgradeAndCall(
                 ITransparentUpgradeableProxy(address(stakingNodesManager)),
-                address(new StakingNodesManager()),
+                address(stakinNodesManagerImplementation),
                 ""
             );
             vm.stopPrank();
@@ -164,6 +203,30 @@ contract Base is Test, Utils {
             stakingNodesManager.upgradeStakingNodeImplementation(address(stakingNodeImplementation));
         }
 
+        // Print StakingNode balance before upgrade
+        console.log("StakingNode balance before upgrade:");
+        for (uint256 i = 0; i < preUpgradeState.stakingNodeBalances.length; i++) {
+            console.log("Node", i, ":", preUpgradeState.stakingNodeBalances[i]);
+        }
+
+        // Print current StakingNode balance after upgrade
+        console.log("StakingNode balance after upgrade:");
+        for (uint256 i = 0; i < stakingNodesManager.nodesLength(); i++) {
+            IStakingNode stakingNode = stakingNodesManager.nodes(i);
+            console.log("Node", i, ":", stakingNode.getETHBalance());
+        }
+
+        // Log pod shares for each eigenpod of each node
+        console.log("EigenPod shares for each StakingNode:");
+        for (uint256 i = 0; i < stakingNodesManager.nodesLength(); i++) {
+            IStakingNode stakingNode = stakingNodesManager.nodes(i);
+            address eigenPodAddress = address(stakingNode.eigenPod());
+            uint256 podShares = uint256(IEigenPodManager(chainAddresses.eigenlayer.EIGENPOD_MANAGER_ADDRESS).podOwnerShares(address(stakingNode)));
+            console.log("Node", i, "Shares:", podShares);
+        }
+
+        //runUpgradeIntegrityInvariants(preUpgradeState);
+
 
         // STAGE 1 - End of atomic upgrade for existing contracts
 
@@ -176,7 +239,6 @@ contract Base is Test, Utils {
             assertGt(previewRedeemAmount, 0, "previewRedeem should return a non-zero value");
         }
 
-
         // STAGE 2 - Deploy and initialize new contracts 
 
         // deploy ynETHRedemptionAssetsVault
@@ -312,7 +374,10 @@ contract Base is Test, Utils {
         for (uint i = 0; i < previousStakingNodeBalances.length; i++) {
             IStakingNode stakingNodeInstance = stakingNodesManager.nodes(i);
             uint256 currentStakingNodeBalance = stakingNodeInstance.getETHBalance();
-            assertEq(currentStakingNodeBalance, previousStakingNodeBalances[i], "Staking node balance integrity check failed for node ID: ");
+            assertEq(
+                currentStakingNodeBalance, previousStakingNodeBalances[i],
+                string.concat("Staking node balance integrity check failed for node ID: ", vm.toString(i))
+            );
         }
 	}
 

diff --git a/test/integration/M3/PlaceholderStakingNodesManager.sol b/test/integration/M3/PlaceholderStakingNodesManager.sol
@@ -0,0 +1,66 @@
+// SPDX-License-Identifier: BSD 3-Clause License
+pragma solidity ^0.8.24;
+
+import {StakingNodesManager} from "src/StakingNodesManager.sol";
+import {IStakingNode} from "src/interfaces/IStakingNode.sol";
+import {IStakingNodesManager} from "src/interfaces/IStakingNodesManager.sol";
+
+contract PlaceholderStakingNodesManager is StakingNodesManager {
+    // Additional functionality can be added here
+
+    struct NodeInitializationData {
+        uint256 initializationDelta;
+    }
+
+    uint256 private immutable node0InitializationDelta;
+    uint256 private immutable node1InitializationDelta;
+    uint256 private immutable node2InitializationDelta;
+    uint256 private immutable node3InitializationDelta;
+    uint256 private immutable node4InitializationDelta;
+
+    constructor(uint256[] memory _nodeDeltas) {
+        require(_nodeDeltas.length == 5, "Invalid number of node deltas");
+
+        node0InitializationDelta = _nodeDeltas[0];
+        node1InitializationDelta = _nodeDeltas[1];
+        node2InitializationDelta = _nodeDeltas[2];
+        node3InitializationDelta = _nodeDeltas[3];
+        node4InitializationDelta = _nodeDeltas[4];
+    }
+
+    function initializeStakingNode(IStakingNode node, uint256 nodeCount) override internal {
+        uint64 initializedVersion = node.getInitializedVersion();
+        if (initializedVersion == 0) {
+            node.initialize(
+                IStakingNode.Init(IStakingNodesManager(address(this)), nodeCount)
+            );
+
+            // update to the newly upgraded version.
+            initializedVersion = node.getInitializedVersion();
+            emit NodeInitialized(address(node), initializedVersion);
+        }
+
+        if (initializedVersion == 1) {
+
+            uint256 initializationDelta;
+            uint256 nodeId = node.nodeId();
+            if (nodeId == 0) {
+                initializationDelta = node0InitializationDelta;
+            } else if (nodeId == 1) {
+                initializationDelta = node1InitializationDelta;
+            } else if (nodeId == 2) {
+                initializationDelta = node2InitializationDelta;
+            } else if (nodeId == 3) {
+                initializationDelta = node3InitializationDelta;
+            } else if (nodeId == 4) {
+                initializationDelta = node4InitializationDelta;
+            } else {
+                initializationDelta = 0;
+            }
+            node.initializeV2(initializationDelta);
+
+        }
+        // NOTE: for future versions add additional if clauses that initialize the node 
+        // for the next version while keeping the previous initializers
+    }
+}
diff --git a/test/integration/M3/WithdrawalsWithRewards-Scenario.t.sol b/test/integration/M3/WithdrawalsWithRewards-Scenario.t.sol
@@ -191,10 +191,7 @@ contract M3WithdrawalsWithRewardsTest is Base {
     }
 
     function test_userWithdrawalWithRewards_Scenario_1() public {
-        // Check if we're on the Holesky testnet
-        if (block.chainid != 17000) {
-            return;
-        }
+
         // deposit 100 ETH into ynETH
         TestState memory state = registerVerifiedValidators(100 ether);
 
@@ -312,10 +309,6 @@ contract M3WithdrawalsWithRewardsTest is Base {
     }
 
     function test_userWithdrawalWithRewards_Scenario_2_queueWithdrawalsBeforeValidatorExit() public {
-        // Check if we're on the Holesky testnet
-        if (block.chainid != 17000) {
-            return;
-        }
 
         // deposit 100 ETH into ynETH
         TestState memory state = registerVerifiedValidators(100 ether);
@@ -436,10 +429,6 @@ contract M3WithdrawalsWithRewardsTest is Base {
 
     function test_userWithdrawalWithRewards_Scenario_3_withdrawEverything() public {
 
-        // Check if we're on the Holesky testnet
-        if (block.chainid != 17000) {
-            return;
-        }
         // exactly 2 validators
         TestState memory state = registerVerifiedValidators(64 ether);
 
@@ -502,10 +491,6 @@ contract M3WithdrawalsWithRewardsTest is Base {
 
     function test_userWithdrawalWithRewards_Scenario_4_slashAllValidatorsAndWithdrawEverything() public {
 
-        // Check if we're on the Holesky testnet
-        if (block.chainid != 17000) {
-            return;
-        }
         // exactly 2 validators
         TestState memory state = registerVerifiedValidators(64 ether);
 
@@ -559,10 +544,6 @@ contract M3WithdrawalsWithRewardsTest is Base {
 
     function test_userWithdrawalWithRewards_Scenario_4_slashAllValidatorsWithNoRewardsAndWithdrawEverything() public {
 
-        // Check if we're on the Holesky testnet
-        if (block.chainid != 17000) {
-            return;
-        }
         // exactly 2 validators
         TestState memory state = registerVerifiedValidators(64 ether);
 
@@ -620,10 +601,6 @@ contract M3WithdrawalsWithRewardsTest is Base {
 
     function test_userWithdrawalWithRewards_Scenario_5_slashAllValidatorsWithNoRewardsAndWithdrawEverything() public {
 
-        // Check if we're on the Holesky testnet
-        if (block.chainid != 17000) {
-            return;
-        }
         // exactly 2 validators
         TestState memory state = registerVerifiedValidators(64 ether);