test: upgrade integration testing

src/test/integration/tests/upgrade/Complete_PreSlashing_Withdrawal.t.sol

@@ -4,104 +4,83 @@ pragma solidity ^0.8.27;
 import "src/test/integration/UpgradeTest.t.sol";
 
 contract Integration_Upgrade_Complete_PreSlashing_Withdrawal is UpgradeTest {
-
-    function testFuzz_deposit_queue_upgrade_completeAsShares(uint24 _random) public rand(_random) {
-        /// Pre-upgrade:
-        /// 1. Create staker with some assets
-        /// 2. Staker deposits into EigenLayer
-        /// 3. Staker queues a withdrawal
-        (User staker, IStrategy[] memory strategies, uint[] memory tokenBalances) = _newRandomStaker();
-        User operator = User(payable(0));
-
-        staker.depositIntoEigenlayer(strategies, tokenBalances);
-        uint[] memory shares = _calculateExpectedShares(strategies, tokenBalances);
-        Withdrawal[] memory withdrawals = staker.queueWithdrawals(strategies, shares);
-
-        /// Upgrade to slashing contracts
-        _upgradeEigenLayerContracts();
-
-        // Complete pre-slashing withdrawals as shares
-        _rollBlocksForCompleteWithdrawals(withdrawals);
-        for (uint i = 0; i < withdrawals.length; i++) {
-            staker.completeWithdrawalAsShares(withdrawals[i]);
-            check_Withdrawal_AsShares_State(staker, operator, withdrawals[i], strategies, shares);
-        }
+    struct TestState {
+        User staker;
+        User operator;
+        IStrategy[] strategies;
+        uint256[] tokenBalances;
+        uint256[] shares;
+        uint256[] withdrawalShares;
+        uint256[] expectedTokens;
+        Withdrawal[] withdrawals;
+        bool isPartial;
+        bool completeAsTokens;
     }
 
-    function testFuzz_delegate_deposit_queue_upgrade_completeAsShares(uint24 _random) public rand(_random) {
-        /// Pre-upgrade:
-        /// 1. Create staker and operator with some asset amounts
-        /// 2. Staker delegates to operator
-        /// 3. Staker deposits into EigenLayer
-        /// 4. Staker queues a withdrawal
-        (User staker, IStrategy[] memory strategies, uint[] memory tokenBalances) = _newRandomStaker();
-        (User operator, ,) = _newRandomOperator();
+    function _init_(uint24 randomness, bool withOperator, bool withDelegation) internal returns (TestState memory state) {
+        (state.staker, state.strategies, state.tokenBalances) = _newRandomStaker();
+        state.shares = _calculateExpectedShares(state.strategies, state.tokenBalances);
 
-        staker.delegateTo(operator);
-        staker.depositIntoEigenlayer(strategies, tokenBalances);
-        uint[] memory shares = _calculateExpectedShares(strategies, tokenBalances);
-        Withdrawal[] memory withdrawals = staker.queueWithdrawals(strategies, shares);
+        state.operator = withOperator ? _newRandomOperator_NoAssets() : User(payable(0));
+        if (withDelegation) state.staker.delegateTo(state.operator);
 
-        /// Upgrade to slashing contracts
-        _upgradeEigenLayerContracts();
+        state.staker.depositIntoEigenlayer(state.strategies, state.tokenBalances);
 
-        // Complete pre-slashing withdrawals as shares
-        _rollBlocksForCompleteWithdrawals(withdrawals);
-        for (uint i = 0; i < withdrawals.length; i++) {
-            staker.completeWithdrawalAsShares(withdrawals[i]);
-            check_Withdrawal_AsShares_State(staker, operator, withdrawals[i], strategies, shares);
+        // Setup withdrawal shares (full or partial)
+        state.isPartial = _randBool();
+        state.withdrawalShares = new uint256[](state.shares.length);
+        for (uint256 i = 0; i < state.shares.length; i++) {
+            state.withdrawalShares[i] = state.isPartial ? state.shares[i] / 2 : state.shares[i];
         }
-    }
 
-    function testFuzz_deposit_queue_upgrade_completeAsTokens(uint24 _random) public rand(_random) {
-        /// Pre-upgrade:
-        /// 1. Create staker with some assets
-        /// 2. Staker deposits into EigenLayer
-        /// 3. Staker queues a withdrawal
-        (User staker, IStrategy[] memory strategies, uint[] memory tokenBalances) = _newRandomStaker();
-        User operator = User(payable(0));
+        state.expectedTokens = _calculateExpectedTokens(state.strategies, state.withdrawalShares);
+        state.completeAsTokens = _randBool();
+    }
 
-        staker.depositIntoEigenlayer(strategies, tokenBalances);
-        uint[] memory shares = _calculateExpectedShares(strategies, tokenBalances);
-        uint[] memory expectedTokens = _calculateExpectedTokens(strategies, shares);
-        Withdrawal[] memory withdrawals = staker.queueWithdrawals(strategies, shares);
-
-        /// Upgrade to slashing contracts
-        _upgradeEigenLayerContracts();
-
-        // Complete pre-slashing withdrawals as tokens
-        _rollBlocksForCompleteWithdrawals(withdrawals);
-        for (uint i = 0; i < withdrawals.length; i++) {
-            IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
-            check_Withdrawal_AsTokens_State(staker, operator, withdrawals[i], strategies, shares, tokens, expectedTokens);
+    function _completeWithdrawal(TestState memory state) internal {
+        for (uint256 i = 0; i < state.withdrawals.length; i++) {
+            if (state.completeAsTokens) {
+                IERC20[] memory tokens = state.staker.completeWithdrawalAsTokens(state.withdrawals[i]);
+                check_Withdrawal_AsTokens_State(state.staker, state.operator, state.withdrawals[i], state.strategies, state.withdrawalShares, tokens, state.expectedTokens);
+            } else {
+                state.staker.completeWithdrawalAsShares(state.withdrawals[i]);
+                check_Withdrawal_AsShares_State(state.staker, state.operator, state.withdrawals[i], state.strategies, state.withdrawalShares);
+            }
         }
     }
 
-    function testFuzz_delegate_deposit_queue_upgrade_completeAsTokens(uint24 _random) public rand(_random) {
-        /// Pre-upgrade:
-        /// 1. Create staker and operator with some asset amounts
-        /// 2. Staker delegates to operator
-        /// 3. Staker deposits into EigenLayer
-        /// 4. Staker queues a withdrawal
-        (User staker, IStrategy[] memory strategies, uint[] memory tokenBalances) = _newRandomStaker();
-        (User operator, ,) = _newRandomOperator();
-
-        staker.delegateTo(operator);
-        staker.depositIntoEigenlayer(strategies, tokenBalances);
-        uint[] memory shares = _calculateExpectedShares(strategies, tokenBalances);
-        uint[] memory expectedTokens = _calculateExpectedTokens(strategies, shares);
-        Withdrawal[] memory withdrawals = staker.queueWithdrawals(strategies, shares);
+    function testFuzz_deposit_queue_upgrade_complete(uint24 r) public rand(r) {
+        TestState memory state = _init_({randomness: r, withOperator: false, withDelegation: false});
+        state.withdrawals = state.staker.queueWithdrawals(state.strategies, state.withdrawalShares);
+        _upgradeEigenLayerContracts();
+        _rollBlocksForCompleteWithdrawals(state.withdrawals);
+        _completeWithdrawal(state);
+    }
 
-        /// Upgrade to slashing contracts
+    function testFuzz_delegate_deposit_queue_upgrade_complete(uint24 r) public rand(r) {
+        TestState memory state = _init_({randomness: r, withOperator: true, withDelegation: true});
+        state.withdrawals = state.staker.queueWithdrawals(state.strategies, state.withdrawalShares);
         _upgradeEigenLayerContracts();
+        _rollBlocksForCompleteWithdrawals(state.withdrawals);
+        _completeWithdrawal(state);
+    }
 
-        // Complete pre-slashing withdrawals as tokens
-        _rollBlocksForCompleteWithdrawals(withdrawals);
-        for (uint i = 0; i < withdrawals.length; i++) {
-            IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
-            check_Withdrawal_AsTokens_State(staker, operator, withdrawals[i], strategies, shares, tokens, expectedTokens);
-        }
+    function testFuzz_upgrade_delegate_queuePartial_complete(uint24 r) public rand(r) {
+        TestState memory state = _init_({randomness: r, withOperator: true, withDelegation: false});
+        _upgradeEigenLayerContracts();
+        state.staker.delegateTo(state.operator);
+        state.withdrawals = state.staker.queueWithdrawals(state.strategies, state.withdrawalShares);
+        _rollBlocksForCompleteWithdrawals(state.withdrawals);
+        _completeWithdrawal(state);
     }
 
-    /// TODO - complete pre-upgrade withdrawal after earliest possible operator slashing
+    function testFuzz_delegate_deposit_queue_completeBeforeUpgrade(uint24 r) public rand(r) {
+        TestState memory state = _init_({randomness: r, withOperator: true, withDelegation: true});
+        state.withdrawals = state.staker.queueWithdrawals(state.strategies, state.withdrawalShares);
+        _rollBlocksForCompleteWithdrawals(state.withdrawals);
+        // We must roll forward by the delay twice since pre-upgrade delay is half as long as post-upgrade delay.
+        rollForward(delegationManager.minWithdrawalDelayBlocks() + 1);
+        _upgradeEigenLayerContracts();
+        _completeWithdrawal(state);
+    }
 }