Integration Tests: Scenarios 2-7

src/test/integration/IntegrationBase.t.sol

@@ -123,18 +123,26 @@ abstract contract IntegrationBase is IntegrationDeployer {
     /// @dev Asserts that ALL of the `withdrawalRoots` is in `delegationManager.pendingWithdrawals`
     function assert_AllWithdrawalsPending(bytes32[] memory withdrawalRoots, string memory err) internal {
         for (uint i = 0; i < withdrawalRoots.length; i++) {
-            assertTrue(delegationManager.pendingWithdrawals(withdrawalRoots[i]), err);
+            assert_WithdrawalPending(withdrawalRoots[i], err);
         }
     }
 
     /// @dev Asserts that NONE of the `withdrawalRoots` is in `delegationManager.pendingWithdrawals`
     function assert_NoWithdrawalsPending(bytes32[] memory withdrawalRoots, string memory err) internal {
         for (uint i = 0; i < withdrawalRoots.length; i++) {
-            assertFalse(delegationManager.pendingWithdrawals(withdrawalRoots[i]), err);
+            assert_WithdrawalNotPending(withdrawalRoots[i], err);
         }
     }
 
     /// @dev Asserts that the hash of each withdrawal corresponds to the provided withdrawal root
+    function assert_WithdrawalPending(bytes32 withdrawalRoot, string memory err) internal {
+        assertTrue(delegationManager.pendingWithdrawals(withdrawalRoot), err);
+    }
+
+    function assert_WithdrawalNotPending(bytes32 withdrawalRoot, string memory err) internal {
+        assertFalse(delegationManager.pendingWithdrawals(withdrawalRoot), err);
+    }
+
     function assert_ValidWithdrawalHashes(
         IDelegationManager.Withdrawal[] memory withdrawals,
         bytes32[] memory withdrawalRoots,
@@ -143,9 +151,17 @@ abstract contract IntegrationBase is IntegrationDeployer {
         bytes32[] memory expectedRoots = _getWithdrawalHashes(withdrawals);
 
         for (uint i = 0; i < withdrawals.length; i++) {
-            assertEq(withdrawalRoots[i], expectedRoots[i], err);
+            assert_ValidWithdrawalHash(withdrawals[i], withdrawalRoots[i], err);
         }
     }
+
+    function assert_ValidWithdrawalHash(
+        IDelegationManager.Withdrawal memory withdrawal,
+        bytes32 withdrawalRoot,
+        string memory err
+    ) internal {
+        assertEq(withdrawalRoot, delegationManager.calculateWithdrawalRoot(withdrawal), err);
+    }
 
     /*******************************************************************************
                                 SNAPSHOT ASSERTIONS
@@ -264,6 +280,45 @@ abstract contract IntegrationBase is IntegrationDeployer {
         }
     }
 
+    function assert_Snap_Removed_StrategyShares(
+        IStrategy[] memory strategies,
+        uint[] memory removedShares,
+        string memory err
+    ) internal {
+        uint[] memory curShares = _getTotalStrategyShares(strategies);
+
+        // Use timewarp to get previous strategy shares
+        uint[] memory prevShares = _getPrevTotalStrategyShares(strategies);
+
+        for (uint i = 0; i < strategies.length; i++) {
+            // Ignore BeaconChainETH strategy since it doesn't keep track of global strategy shares
+            if (strategies[i] == BEACONCHAIN_ETH_STRAT) {
+                continue;
+            }
+            uint prevShare = prevShares[i];
+            uint curShare = curShares[i];
+
+            assertEq(prevShare - removedShares[i], curShare, err);
+        }
+    }
+
+    function assert_Snap_Unchanged_StrategyShares(
+        IStrategy[] memory strategies,
+        string memory err
+    ) internal {
+        uint[] memory curShares = _getTotalStrategyShares(strategies);
+
+        // Use timewarp to get previous strategy shares
+        uint[] memory prevShares = _getPrevTotalStrategyShares(strategies);
+
+        for (uint i = 0; i < strategies.length; i++) {
+            uint prevShare = prevShares[i];
+            uint curShare = curShares[i];
+
+            assertEq(prevShare, curShare, err);
+        }
+    }
+
     /// Snapshot assertions for underlying token balances:
 
     /// @dev Check that the staker has `addedTokens` additional underlying tokens 
@@ -339,6 +394,18 @@ abstract contract IntegrationBase is IntegrationDeployer {
         assertEq(prevQueuedWithdrawals + withdrawals.length, curQueuedWithdrawals, err);
     }
 
+    function assert_Snap_Added_QueuedWithdrawal(
+        User staker, 
+        IDelegationManager.Withdrawal memory withdrawal,
+        string memory err
+    ) internal {
+        uint curQueuedWithdrawal = _getCumulativeWithdrawals(staker);
+        // Use timewarp to get previous cumulative withdrawals
+        uint prevQueuedWithdrawal = _getPrevCumulativeWithdrawals(staker);
+
+        assertEq(prevQueuedWithdrawal + 1, curQueuedWithdrawal, err);
+    }
+
     /*******************************************************************************
                                 UTILITY METHODS
     *******************************************************************************/
@@ -375,6 +442,10 @@ abstract contract IntegrationBase is IntegrationDeployer {
         return (withdrawStrats, withdrawShares);
     }
 
+    /**
+     * Helpful getters:
+     */
+
     /// @dev For some strategies/underlying token balances, calculate the expected shares received
     /// from depositing all tokens
     function _calculateExpectedShares(IStrategy[] memory strategies, uint[] memory tokenBalances) internal returns (uint[] memory) {
@@ -524,4 +595,21 @@ abstract contract IntegrationBase is IntegrationDeployer {
 
         return balances;
     }
+
+    function _getPrevTotalStrategyShares(IStrategy[] memory strategies) internal timewarp() returns (uint[] memory) {
+        return _getTotalStrategyShares(strategies);
+    }
+
+    function _getTotalStrategyShares(IStrategy[] memory strategies) internal view returns (uint[] memory) {
+        uint[] memory shares = new uint[](strategies.length);
+
+        for (uint i = 0; i < strategies.length; i++) {
+            if (strategies[i] != BEACONCHAIN_ETH_STRAT) {
+                shares[i] = strategies[i].totalShares();
+            }
+            // BeaconChainETH strategy doesn't keep track of global strategy shares, so we ignore
+        }
+
+        return shares;
+    }
 }

src/test/integration/IntegrationChecks.t.sol

@@ -0,0 +1,130 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity =0.8.12;
+
+import "src/test/integration/IntegrationBase.t.sol";
+import "src/test/integration/User.t.sol";
+
+/// @notice Contract that provides utility functions to reuse common test blocks & checks
+contract IntegrationCheckUtils is IntegrationBase {
+
+    function check_Deposit_State(User staker, IStrategy[] memory strategies, uint[] memory shares) internal {
+        /// Deposit into strategies:
+        // For each of the assets held by the staker (either StrategyManager or EigenPodManager),
+        // the staker calls the relevant deposit function, depositing all held assets.
+        //
+        // ... check that all underlying tokens were transferred to the correct destination
+        //     and that the staker now has the expected amount of delegated shares in each strategy
+        assert_HasNoUnderlyingTokenBalance(staker, strategies, "staker should have transferred all underlying tokens");
+        assert_Snap_Added_StakerShares(staker, strategies, shares, "staker should expected shares in each strategy after depositing");
+    }
+
+    function check_Delegation_State(User staker, User operator, IStrategy[] memory strategies, uint[] memory shares) internal {
+        /// Delegate to an operator:
+        //
+        // ... check that the staker is now delegated to the operator, and that the operator
+        //     was awarded the staker shares
+        assertTrue(delegationManager.isDelegated(address(staker)), "staker should be delegated");
+        assertEq(address(operator), delegationManager.delegatedTo(address(staker)), "staker should be delegated to operator");
+        assert_HasExpectedShares(staker, strategies, shares, "staker should still have expected shares after delegating");
+        assert_Snap_Added_OperatorShares(operator, strategies, shares, "operator should have received shares");
+    }
+
+    function check_QueuedWithdrawal_State(User staker, User operator, IStrategy[] memory strategies, uint[] memory shares, IDelegationManager.Withdrawal[] memory withdrawals, bytes32[] memory withdrawalRoots) internal {
+        // The staker will queue one or more withdrawals for the selected strategies and shares
+        //
+        // ... check that each withdrawal was successfully enqueued, that the returned roots
+        //     match the hashes of each withdrawal, and that the staker and operator have
+        //     reduced shares.
+        assert_AllWithdrawalsPending(withdrawalRoots, "staker withdrawals should now be pending");
+        assert_ValidWithdrawalHashes(withdrawals, withdrawalRoots, "calculated withdrawals should match returned roots");
+        assert_Snap_Added_QueuedWithdrawals(staker, withdrawals, "staker should have increased nonce by withdrawals.length");
+        assert_Snap_Removed_OperatorShares(operator, strategies, shares, "failed to remove operator shares");
+        assert_Snap_Removed_StakerShares(staker, strategies, shares, "failed to remove staker shares");
+    }
+
+    function check_Undelegate_State(
+        User staker, 
+        User operator, 
+        IDelegationManager.Withdrawal[] memory withdrawals,
+        bytes32[] memory withdrawalRoots,
+        IStrategy[] memory strategies,
+        uint[] memory shares 
+    ) internal {
+        /// Undelegate from an operator
+        //
+        // ... check that the staker is undelegated, all strategies from which the staker is deposited are unqeuued,
+        //     that the returned root matches the hashes for each strategy and share amounts, and that the staker
+        //     and operator have reduced shares
+        assertEq(withdrawalRoots.length, 1, "should only be one withdrawal root");
+        assertFalse(delegationManager.isDelegated(address(staker)), "staker should not be delegated");
+        assert_ValidWithdrawalHashes(withdrawals, withdrawalRoots, "calculated withdrawl should match returned root");
+        assert_AllWithdrawalsPending(withdrawalRoots, "stakers withdrawal should now be pending");
+        assert_Snap_Added_QueuedWithdrawals(staker, withdrawals, "staker should have increased nonce by 1");
+        assert_Snap_Removed_OperatorShares(operator, strategies, shares, "failed to remove operator shares");
+        assert_Snap_Removed_StakerShares(staker, strategies, shares, "failed to remove staker shares");
+    }
+
+    function check_Withdrawal_AsTokens_State(
+        User staker,
+        User operator,
+        IDelegationManager.Withdrawal memory withdrawal,
+        IStrategy[] memory strategies,
+        uint[] memory shares,
+        IERC20[] memory tokens,
+        uint[] memory expectedTokens
+    ) internal {
+        /// Complete withdrawal(s):
+        // The staker will complete the withdrawal as tokens
+        // 
+        // ... check that the withdrawal is not pending, that the withdrawer received the expected tokens, and that the total shares of each 
+        //     strategy withdrawn decreases
+        assert_WithdrawalNotPending(delegationManager.calculateWithdrawalRoot(withdrawal), "staker withdrawal should no longer be pending");
+        assert_Snap_Added_TokenBalances(staker, tokens, expectedTokens, "staker should have received expected tokens");
+        assert_Snap_Unchanged_TokenBalances(operator, "operator token balances should not have changed");
+        assert_Snap_Unchanged_StakerShares(staker, "staker shares should not have changed");
+        assert_Snap_Unchanged_OperatorShares(operator, "operator shares should not have changed");
+        assert_Snap_Removed_StrategyShares(strategies, shares, "strategies should have total shares decremented");
+    }
+
+    function check_Withdrawal_AsShares_State(
+        User staker,
+        User operator,
+        IDelegationManager.Withdrawal memory withdrawal,
+        IStrategy[] memory strategies,
+        uint[] memory shares
+    ) internal {
+        /// Complete withdrawal(s):
+        // The staker will complete the withdrawal as shares
+        // 
+        // ... check that the withdrawal is not pending, that the withdrawer received the expected shares, and that the total shares of each 
+        //     strategy withdrawn remains unchanged 
+        assert_WithdrawalNotPending(delegationManager.calculateWithdrawalRoot(withdrawal), "staker withdrawal should no longer be pending");
+        assert_Snap_Unchanged_TokenBalances(staker, "staker should not have any change in underlying token balances");
+        assert_Snap_Unchanged_TokenBalances(operator, "operator should not have any change in underlying token balances");
+        assert_Snap_Added_StakerShares(staker, strategies, shares, "staker should have received expected shares");
+        assert_Snap_Added_OperatorShares(operator, withdrawal.strategies, withdrawal.shares, "operator should have received shares");
+        assert_Snap_Unchanged_StrategyShares(strategies, "strategies should have total shares unchanged");
+    }
+
+    /// @notice Difference from above is that operator shares do not increase since staker is not delegated
+    function check_Withdrawal_AsShares_Undelegated_State(
+        User staker,
+        User operator,
+        IDelegationManager.Withdrawal memory withdrawal,
+        IStrategy[] memory strategies,
+        uint[] memory shares
+    ) internal {
+        /// Complete withdrawal(s):
+        // The staker will complete the withdrawal as shares
+        // 
+        // ... check that the withdrawal is not pending, that the token balances of the staker and operator are unchanged,
+        //     that the withdrawer received the expected shares, and that that the total shares of each o
+        //     strategy withdrawn remains unchanged 
+        assert_WithdrawalNotPending(delegationManager.calculateWithdrawalRoot(withdrawal), "staker withdrawal should no longer be pending");
+        assert_Snap_Unchanged_TokenBalances(staker, "staker should not have any change in underlying token balances");
+        assert_Snap_Unchanged_TokenBalances(operator, "operator should not have any change in underlying token balances");
+        assert_Snap_Added_StakerShares(staker, strategies, shares, "staker should have received expected shares");
+        assert_Snap_Unchanged_OperatorShares(operator, "operator should have shares unchanged");
+        assert_Snap_Unchanged_StrategyShares(strategies, "strategies should have total shares unchanged");
+    }
+}

src/test/integration/IntegrationDeployer.t.sol

@@ -231,7 +231,7 @@ abstract contract IntegrationDeployer is Test, IUserDeployer {
             address(eigenPodManagerImplementation),
             abi.encodeWithSelector(
                 EigenPodManager.initialize.selector,
-                type(uint256).max, // maxPods
+                type(uint).max, // maxPods
                 address(beaconChainOracle),
                 eigenLayerReputedMultisig, // initialOwner
                 pauserRegistry,

src/test/integration/User.t.sol

@@ -123,6 +123,22 @@ contract User is Test {
         delegationManager.delegateTo(address(operator), emptySig, bytes32(0));
     }
 
+    /// @dev Undelegate from operator
+    function undelegate() public createSnapshot virtual returns(IDelegationManager.Withdrawal[] memory){
+        IDelegationManager.Withdrawal[] memory withdrawal = new IDelegationManager.Withdrawal[](1);
+        withdrawal[0] = _getExpectedWithdrawalStructForStaker(address(this));
+        delegationManager.undelegate(address(this));
+        return withdrawal;
+    }
+
+    /// @dev Force undelegate staker
+    function forceUndelegate(User staker) public createSnapshot virtual returns(IDelegationManager.Withdrawal[] memory){
+        IDelegationManager.Withdrawal[] memory withdrawal = new IDelegationManager.Withdrawal[](1);
+        withdrawal[0] = _getExpectedWithdrawalStructForStaker(address(staker));
+        delegationManager.undelegate(address(staker));
+        return withdrawal;
+    }
+
     /// @dev Queues a single withdrawal for every share and strategy pair
     function queueWithdrawals(
         IStrategy[] memory strategies, 
@@ -160,11 +176,19 @@ contract User is Test {
 
         return (withdrawals);
     }
+
+    function completeWithdrawalAsTokens(IDelegationManager.Withdrawal memory withdrawal) public createSnapshot virtual returns (IERC20[] memory) {
+        return _completeQueuedWithdrawal(withdrawal, true);
+    }
+
+    function completeWithdrawalAsShares(IDelegationManager.Withdrawal memory withdrawal) public createSnapshot virtual returns (IERC20[] memory) {
+        return _completeQueuedWithdrawal(withdrawal, false);
+    }
 
-    function completeQueuedWithdrawal(
+    function _completeQueuedWithdrawal(
         IDelegationManager.Withdrawal memory withdrawal, 
         bool receiveAsTokens
-    ) public createSnapshot virtual returns (IERC20[] memory) {
+    ) internal virtual returns (IERC20[] memory) {
         IERC20[] memory tokens = new IERC20[](withdrawal.strategies.length);
 
         for (uint i = 0; i < tokens.length; i++) {
@@ -215,6 +239,22 @@ contract User is Test {
     function _podWithdrawalCredentials() internal view returns (bytes memory) {
         return abi.encodePacked(bytes1(uint8(1)), bytes11(0), address(pod));
     }
+
+    /// @notice Assumes staker and withdrawer are the same and that all strategies and shares are withdrawn
+    function _getExpectedWithdrawalStructForStaker(address staker) internal view returns (IDelegationManager.Withdrawal memory) {
+        (IStrategy[] memory strategies, uint[] memory shares)
+            = delegationManager.getDelegatableShares(staker);
+
+        return IDelegationManager.Withdrawal({
+            staker: staker,
+            delegatedTo: delegationManager.delegatedTo(staker),
+            withdrawer: staker,
+            nonce: delegationManager.cumulativeWithdrawalsQueued(staker),
+            startBlock: uint32(block.number),
+            strategies: strategies,
+            shares: shares
+        });
+    }
 }
 
 /// @notice A user contract that calls nonstandard methods (like xBySignature methods)