test: integration w/ multiple stakers

8sunyuan · 8sunyuan · commit 4cd8b1316904 · 2025-02-19T17:55:09.000-05:00
diff --git a/src/test/integration/IntegrationBase.t.sol b/src/test/integration/IntegrationBase.t.sol
@@ -69,6 +69,23 @@ abstract contract IntegrationBase is IntegrationDeployer, TypeImporter {
         return (staker, strategies, tokenBalances);
     }
 
+    /// @dev creates `numNewStakers` new stakers with random strategies and token balances
+    function _newRandomStakers(uint numNewStakers) internal returns (User[] memory, IStrategy[][] memory, uint[][] memory) {
+        User[] memory stakers = new User[](numNewStakers);
+        IStrategy[][] memory strategies = new IStrategy[][](numNewStakers);
+        uint[][] memory tokenBalances = new uint[][](numNewStakers);
+
+        for (uint i = 0; i < numNewStakers; i++) {
+            (stakers[i], strategies[i], tokenBalances[i]) = _newRandomStaker();
+        }
+
+        numStakers += numNewStakers;
+    
+        return (stakers, strategies, tokenBalances);
+    }
+
+    /// @dev Exactly like _newRandomStaker above but sets strategies and tokenBalances arrays to length 1
+    /// to use a single strategy and balance
     function _newBasicStaker() internal returns (User, IStrategy[] memory, uint[] memory) {
         string memory stakerName;
 
@@ -2307,6 +2324,21 @@ abstract contract IntegrationBase is IntegrationDeployer, TypeImporter {
         return expectedShares;
     }
 
+    /// @dev For some strategies/underlying token balances, calculate the expected shares received
+    /// Exact same as _calculateExpectedShares but for multiple stakers
+    function _calculateExpectedShares(
+        IStrategy[][] memory strategies,
+        uint[][] memory tokenBalances
+    ) internal returns (uint[][] memory) {
+        uint[][] memory expectedShares = new uint[][](strategies.length);
+
+        for (uint i = 0; i < strategies.length; i++) {
+            expectedShares[i] = _calculateExpectedShares(strategies[i], tokenBalances[i]);
+        }
+
+        return expectedShares;
+    }
+
     /// @dev For some strategies/underlying token balances, calculate the expected shares received
     /// from depositing all tokens
     function _calculateExpectedTokens(IStrategy[] memory strategies, uint[] memory shares) internal returns (uint[] memory) {
diff --git a/src/test/integration/tests/multiStaker/Deposit_Delegate_Queue_Complete.t.sol b/src/test/integration/tests/multiStaker/Deposit_Delegate_Queue_Complete.t.sol
@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity ^0.8.27;
+
+import "src/test/integration/IntegrationChecks.t.sol";
+import "src/test/integration/users/User.t.sol";
+
+import "src/test/integration/utils/MultiStakersLib.sol";
+
+contract Integration_MultiStaker_Deposit_Delegate_Queue_Complete is IntegrationCheckUtils {
+    using MultiStakersLib for User[];
+
+    /*******************************************************************************
+                                FULL WITHDRAWALS
+    *******************************************************************************/
+
+    // TODO: fix test
+    /// Generates multiple random stakers and an operator. The stakers:
+    /// 1. deposits all assets into strategies
+    /// 2. delegates to an operator
+    /// 3. queues a withdrawal for ALL shares
+    /// 4. completes the queued withdrawal as tokens
+    function testFuzz_deposit_delegate_queue_completeAsTokens(uint24 _random) public rand(_random) {   
+        /// 0. Create multiple random stakers with:
+        // - some nonzero underlying token balances
+        // - corresponding to a random subset of valid strategies (StrategyManager and/or EigenPodManager)
+        //
+        // ... check that the staker has no delegatable shares and isn't currently delegated
+
+        ( 
+            User[] memory stakers,
+            IStrategy[][] memory strategiesPerStaker,
+            uint[][] memory tokenBalancesPerStaker
+        ) = _newRandomStakers({numNewStakers: _randUint(1, 10)});
+        (User operator, ,) = _newRandomOperator();
+
+        uint[][] memory sharesPerStaker = _calculateExpectedShares(strategiesPerStaker, tokenBalancesPerStaker);
+
+        // 1. Deposit Into Strategies
+        stakers.depositIntoEigenlayer(strategiesPerStaker, tokenBalancesPerStaker);
+        // check_Deposit_State(staker, strategies, shares);
+
+        // 2. Delegate to an operator
+        stakers.delegateTo(operator);
+        // check_Delegation_State(staker, operator, strategies, shares);
+
+        // 3. Queue Withdrawals
+        Withdrawal[][] memory withdrawals = stakers.queueWithdrawals(strategiesPerStaker, sharesPerStaker);
+
+        // 4. Complete withdrawal
+        _rollBlocksForCompleteWithdrawals(withdrawals[0]);
+        stakers.completeWithdrawalsAsTokens(withdrawals);
+    }
+}
diff --git a/src/test/integration/users/User.t.sol b/src/test/integration/users/User.t.sol
@@ -34,6 +34,7 @@ contract User is Logger, IDelegationManagerTypes, IAllocationManagerTypes {
     using SlashingLib for *;
     using ArrayLib for *;
     using print for *;
+    using EnumerableSet for EnumerableSet.AddressSet;
 
     IStrategy constant beaconChainETHStrategy = IStrategy(0xbeaC0eeEeeeeEEeEeEEEEeeEEeEeeeEeeEEBEaC0);
 
@@ -52,6 +53,10 @@ contract User is Logger, IDelegationManagerTypes, IAllocationManagerTypes {
 
     string _NAME;
 
+    // User's delegated stakers if they are an operator
+    // create an enumerableSet of stakers that have delegated to this user
+    EnumerableSet.AddressSet internal delegatedStakers;
+
     // User's EigenPod and each of their validator indices within that pod
     EigenPod public pod;
     uint40[] validators;
@@ -208,11 +213,16 @@ contract User is Logger, IDelegationManagerTypes, IAllocationManagerTypes {
         ISignatureUtilsMixinTypes.SignatureWithExpiry memory emptySig;
         delegationManager.delegateTo(address(operator), emptySig, bytes32(0));
         print.gasUsed();
+
+        // handle the delegating staker on the operator User contract
+        // to manage the delegatedStakers enum set
+        operator.handleDelegatingStakers(address(this));
     }
 
     /// @dev Undelegate from operator
     function undelegate() public virtual createSnapshot returns (Withdrawal[] memory) {
         print.method("undelegate");
+        User operator = User(payable(delegationManager.delegatedTo(address(this))));
 
         Withdrawal[] memory expectedWithdrawals = _getExpectedWithdrawalStructsForStaker(address(this));
         _tryPrankAppointee_DelegationManager(IDelegationManager.undelegate.selector);
@@ -234,6 +244,10 @@ contract User is Logger, IDelegationManagerTypes, IAllocationManagerTypes {
             );
         }
 
+        // handle the undelegating staker on the operator User contract
+        // to manage the delegatedStakers enum set
+        operator.handleUndelegatingStaker(address(this));
+
         return expectedWithdrawals;
     }
 
@@ -242,6 +256,8 @@ contract User is Logger, IDelegationManagerTypes, IAllocationManagerTypes {
         User newOperator
     ) public virtual createSnapshot returns (Withdrawal[] memory) {
         print.method("redelegate", newOperator.NAME_COLORED());
+        User oldOperator = User(payable(delegationManager.delegatedTo(payable(address(this)))));
+
         Withdrawal[] memory expectedWithdrawals = _getExpectedWithdrawalStructsForStaker(address(this));
         ISignatureUtilsMixinTypes.SignatureWithExpiry memory emptySig;
         _tryPrankAppointee_DelegationManager(IDelegationManager.redelegate.selector);
@@ -262,6 +278,12 @@ contract User is Logger, IDelegationManagerTypes, IAllocationManagerTypes {
                 expectedWithdrawals[i].scaledShares[0]
             );
         }
+
+        // handle the redelegating staker on the operator User contract
+        // to manage the delegatedStakers enum set
+        oldOperator.handleUndelegatingStaker(address(this));
+        newOperator.handleDelegatingStakers(address(this));
+
         return expectedWithdrawals;
     }
 
@@ -275,9 +297,35 @@ contract User is Logger, IDelegationManagerTypes, IAllocationManagerTypes {
         delegationManager.undelegate(address(staker));
         print.gasUsed();
 
+        // Operator is calling this function themselves
+        // handle the undelegating staker on this contract to manage the delegatedStakers enum set
+        handleUndelegatingStaker(address(staker));
+
         return expectedWithdrawals;
     }
 
+    /// @dev NOTE: do NOT call this function directly. This is simply implemented to manage the
+    /// delegatedStakers enum set. This function is only meant to be called on the Operator User contract.
+    function handleDelegatingStakers(address staker) public virtual {
+        // check that the current User is an operator
+        require(
+            delegationManager.isOperator(address(this)),
+            "User is not an operator"
+        );
+        delegatedStakers.add(staker);
+    }
+
+    /// @dev NOTE: do NOT call this function directly. This is simply implemented to manage the
+    /// delegatedStakers enum set. This function is only meant to be called on the Operator User contract.
+    function handleUndelegatingStaker(address staker) public virtual {
+        // check that the current User is an operator
+        require(
+            delegationManager.isOperator(address(this)),
+            "User is not an operator"
+        );
+        delegatedStakers.remove(staker);
+    }
+
     /// @dev Queues a single withdrawal for every share and strategy pair
     function queueWithdrawals(
         IStrategy[] memory strategies,
@@ -713,6 +761,15 @@ contract User is Logger, IDelegationManagerTypes, IAllocationManagerTypes {
         return activeValidators;
     }
 
+    function getDelegatedStakers() public view returns (User[] memory) {
+        address[] memory stakers = delegatedStakers.values();
+        User[] memory users = new User[](stakers.length);
+        for (uint256 i = 0; i < stakers.length; i++) {
+            users[i] = User(payable(stakers[i]));
+        }
+        return users;
+    }
+
     function _tryPrankAppointee(
         address target,
         bytes4 selector
diff --git a/src/test/integration/utils/MultiStakersLib.sol b/src/test/integration/utils/MultiStakersLib.sol
@@ -0,0 +1,75 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity ^0.8.27;
+
+import "src/test/integration/users/User.t.sol";
+
+
+/// @notice Contract that provides utility functions to reuse staker actions in User.t.sol
+// for multiple stakers at once. 
+library MultiStakersLib {
+
+    function depositIntoEigenlayer(
+        User[] memory stakers,
+        IStrategy[][] memory strategies,
+        uint[][] memory tokenBalances
+    ) internal {
+        for (uint i = 0; i < stakers.length; i++) {
+            stakers[i].depositIntoEigenlayer(strategies[i], tokenBalances[i]);
+        }
+    }
+
+    function delegateTo(
+        User[] memory stakers,
+        User operator
+    ) internal {
+        for (uint i = 0; i < stakers.length; i++) {
+            stakers[i].delegateTo(operator);
+        }
+    }
+
+    function undelegate(
+        User[] memory stakers
+    ) internal {
+        for (uint i = 0; i < stakers.length; i++) {
+            stakers[i].undelegate();
+        }
+    }
+
+    function redelegate(
+        User[] memory stakers,
+        User newOperator
+    ) internal {
+        for (uint i = 0; i < stakers.length; i++) {
+            stakers[i].redelegate(newOperator);
+        }
+    }
+
+    function queueWithdrawals(
+        User[] memory stakers,
+        IStrategy[][] memory strategies,
+        uint[][] memory shares
+    ) internal returns (IDelegationManagerTypes.Withdrawal[][] memory withdrawals) {
+        withdrawals = new IDelegationManagerTypes.Withdrawal[][](stakers.length);
+        for (uint i = 0; i < stakers.length; i++) {
+            withdrawals[i] = stakers[i].queueWithdrawals(strategies[i], shares[i]);
+        }
+    }
+
+    function completeWithdrawalsAsTokens(
+        User[] memory stakers,
+        IDelegationManagerTypes.Withdrawal[][] memory withdrawals
+    ) internal {
+        for (uint i = 0; i < stakers.length; i++) {
+            stakers[i].completeWithdrawalsAsTokens(withdrawals[i]);
+        }
+    }
+
+    function completeWithdrawalsAsShares(
+        User[] memory stakers,
+        IDelegationManagerTypes.Withdrawal[][] memory withdrawals
+    ) internal {
+        for (uint i = 0; i < stakers.length; i++) {
+            stakers[i].completeWithdrawalsAsShares(withdrawals[i]);
+        }
+    }
+}
