chore: fix tests to work with modified behavior (#378)

ChaoticWalrus · wadealexc · web-flow · commit 71e82eea0abf · 2023-12-17T17:44:38.000-05:00
* chore: fix tests to work with modified behavior

integration tests in particular are now slightly more flexible

* fix: remove memory overwrite

* docs: update dmgr docs

---------

Co-authored-by: wadealexc &lt;pragma-services@proton.me&gt;
diff --git a/docs/core/DelegationManager.md b/docs/core/DelegationManager.md
@@ -229,24 +229,24 @@ function undelegate(
 ) 
     external 
     onlyWhenNotPaused(PAUSED_ENTER_WITHDRAWAL_QUEUE)
-    returns (bytes32 withdrawalRoot)
+    returns (bytes32[] memory withdrawalRoots)
 ```
 
-`undelegate` can be called by a Staker to undelegate themselves, or by a Staker's delegated Operator (or that Operator's `delegationApprover`). Undelegation (i) queues a withdrawal on behalf of the Staker for all their delegated shares, and (ii) decreases the Operator's delegated shares according to the amounts and strategies being withdrawn.
+`undelegate` can be called by a Staker to undelegate themselves, or by a Staker's delegated Operator (or that Operator's `delegationApprover`). Undelegation (i) queues withdrawals on behalf of the Staker for all their delegated shares, and (ii) decreases the Operator's delegated shares according to the amounts and strategies being withdrawn.
 
-If the Staker has active shares in either the `EigenPodManager` or `StrategyManager`, they are removed while the withdrawal is in the queue.
+If the Staker has active shares in either the `EigenPodManager` or `StrategyManager`, they are removed while the withdrawal is in the queue - and an individual withdrawal is queued for each strategy removed.
 
-The withdrawal can be completed by the Staker after `withdrawalDelayBlocks`, and does not require the Staker to "fully exit" from the system -- the Staker may choose to receive their shares back in full once the withdrawal is completed (see [`completeQueuedWithdrawal`](#completequeuedwithdrawal) for details).
+The withdrawals can be completed by the Staker after `withdrawalDelayBlocks`. This does not require the Staker to "fully exit" from the system -- the Staker may choose to receive their shares back in full once withdrawals are completed (see [`completeQueuedWithdrawal`](#completequeuedwithdrawal) for details).
 
 Note that becoming an Operator is irreversible! Although Operators can withdraw, they cannot use this method to undelegate from themselves.
 
 *Effects*: 
 * Any shares held by the Staker in the `EigenPodManager` and `StrategyManager` are removed from the Operator's delegated shares.
 * The Staker is undelegated from the Operator
 * If the Staker has no delegatable shares, there is no withdrawal queued or further effects
-* A `Withdrawal` is queued for the Staker, tracking the strategies and shares being withdrawn
-    * The Staker's withdrawal nonce is increased
-    * The hash of the `Withdrawal` is marked as "pending"
+* For each strategy being withdrawn, a `Withdrawal` is queued for the Staker:
+    * The Staker's withdrawal nonce is increased by 1 for each `Withdrawal`
+    * The hash of each `Withdrawal` is marked as "pending"
 * See [`EigenPodManager.removeShares`](./EigenPodManager.md#eigenpodmanagerremoveshares)
 * See [`StrategyManager.removeShares`](./StrategyManager.md#removeshares)
 
diff --git a/src/contracts/core/DelegationManager.sol b/src/contracts/core/DelegationManager.sol
@@ -242,12 +242,12 @@ contract DelegationManager is Initializable, OwnableUpgradeable, Pausable, Deleg
         emit StakerUndelegated(staker, operator);
         delegatedTo[staker] = address(0);
 
-        // if no delegatable shares, return zero root, and don't queue a withdrawal
+        // if no delegatable shares, return an empty array, and don't queue a withdrawal
         if (strategies.length == 0) {
             withdrawalRoots = new bytes32[](0);
         } else {
             withdrawalRoots = new bytes32[](strategies.length);
-            for (uint i = 0; i < strategies.length; i++) {
+            for (uint256 i = 0; i < strategies.length; i++) {
                 IStrategy[] memory singleStrategy = new IStrategy[](1);
                 uint256[] memory singleShare = new uint256[](1);
                 singleStrategy[0] = strategies[i];
diff --git a/src/test/integration/IntegrationChecks.t.sol b/src/test/integration/IntegrationChecks.t.sol
@@ -52,11 +52,16 @@ contract IntegrationCheckUtils is IntegrationBase {
         // ... 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");
+        assert_AllWithdrawalsPending(withdrawalRoots,
+            "check_QueuedWithdrawal_State: staker withdrawals should now be pending");
+        assert_ValidWithdrawalHashes(withdrawals, withdrawalRoots,
+            "check_QueuedWithdrawal_State: calculated withdrawals should match returned roots");
+        assert_Snap_Added_QueuedWithdrawals(staker, withdrawals,
+            "check_QueuedWithdrawal_State: staker should have increased nonce by withdrawals.length");
+        assert_Snap_Removed_OperatorShares(operator, strategies, shares,
+            "check_QueuedWithdrawal_State: failed to remove operator shares");
+        assert_Snap_Removed_StakerShares(staker, strategies, shares,
+            "check_QueuedWithdrawal_State: failed to remove staker shares");
     }
 
     function check_Undelegate_State(
@@ -72,13 +77,18 @@ contract IntegrationCheckUtils is IntegrationBase {
         // ... 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");
+        assertFalse(delegationManager.isDelegated(address(staker)),
+            "check_Undelegate_State: staker should not be delegated");
+        assert_ValidWithdrawalHashes(withdrawals, withdrawalRoots,
+            "check_Undelegate_State: calculated withdrawl should match returned root");
+        assert_AllWithdrawalsPending(withdrawalRoots,
+            "check_Undelegate_State: stakers withdrawal should now be pending");
+        assert_Snap_Added_QueuedWithdrawals(staker, withdrawals,
+            "check_Undelegate_State: staker should have increased nonce by withdrawals.length");
+        assert_Snap_Removed_OperatorShares(operator, strategies, shares,
+            "check_Undelegate_State: failed to remove operator shares");
+        assert_Snap_Removed_StakerShares(staker, strategies, shares,
+            "check_Undelegate_State: failed to remove staker shares");
     }
 
     function check_Withdrawal_AsTokens_State(
diff --git a/src/test/integration/User.t.sol b/src/test/integration/User.t.sol
@@ -171,20 +171,26 @@ contract User is Test {
     function undelegate() public createSnapshot virtual returns(IDelegationManager.Withdrawal[] memory){
         emit log(_name(".undelegate"));
 
-        IDelegationManager.Withdrawal[] memory withdrawal = new IDelegationManager.Withdrawal[](1);
-        withdrawal[0] = _getExpectedWithdrawalStructForStaker(address(this));
+        IDelegationManager.Withdrawal[] memory expectedWithdrawals = _getExpectedWithdrawalStructsForStaker(address(this));
         delegationManager.undelegate(address(this));
-        return withdrawal;
+
+        for (uint i = 0; i < expectedWithdrawals.length; i++) {
+            emit log("expecting withdrawal:");
+            emit log_named_uint("nonce: ", expectedWithdrawals[i].nonce);
+            emit log_named_address("strat: ", address(expectedWithdrawals[i].strategies[0]));
+            emit log_named_uint("shares: ", expectedWithdrawals[i].shares[0]);
+        }
+        
+        return expectedWithdrawals;
     }
 
     /// @dev Force undelegate staker
     function forceUndelegate(User staker) public createSnapshot virtual returns(IDelegationManager.Withdrawal[] memory){
         emit log_named_string(_name(".forceUndelegate: "), staker.NAME());
 
-        IDelegationManager.Withdrawal[] memory withdrawal = new IDelegationManager.Withdrawal[](1);
-        withdrawal[0] = _getExpectedWithdrawalStructForStaker(address(staker));
+        IDelegationManager.Withdrawal[] memory expectedWithdrawals = _getExpectedWithdrawalStructsForStaker(address(staker));
         delegationManager.undelegate(address(staker));
-        return withdrawal;
+        return expectedWithdrawals;
     }
 
     /// @dev Queues a single withdrawal for every share and strategy pair
@@ -317,20 +323,33 @@ contract User is Test {
         return abi.encodePacked(bytes1(uint8(1)), bytes11(0), address(pod));
     }
 
+    /// @notice Gets the expected withdrawals to be created when the staker is undelegated via a call to `DelegationManager.undelegate()`
     /// @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)
+    function _getExpectedWithdrawalStructsForStaker(address staker) internal returns (IDelegationManager.Withdrawal[] memory) {
+        (IStrategy[] memory strategies, uint256[] 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
-        });
+        IDelegationManager.Withdrawal[] memory expectedWithdrawals = new IDelegationManager.Withdrawal[](strategies.length);
+        address delegatedTo = delegationManager.delegatedTo(staker);
+        uint256 nonce = delegationManager.cumulativeWithdrawalsQueued(staker);
+        
+        for (uint256 i = 0; i < strategies.length; ++i) {
+            IStrategy[] memory singleStrategy = new IStrategy[](1);
+            uint256[] memory singleShares = new uint256[](1);
+            singleStrategy[0] = strategies[i];
+            singleShares[0] = shares[i];
+            expectedWithdrawals[i] = IDelegationManager.Withdrawal({
+                staker: staker,
+                delegatedTo: delegatedTo,
+                withdrawer: staker,
+                nonce: (nonce + i),
+                startBlock: uint32(block.number),
+                strategies: singleStrategy,
+                shares: singleShares
+            });
+        }
+
+        return expectedWithdrawals;
     }
 
     function _name(string memory s) internal view returns (string memory) {
diff --git a/src/test/integration/tests/Deposit_Delegate_Queue_Complete.t.sol b/src/test/integration/tests/Deposit_Delegate_Queue_Complete.t.sol
@@ -56,8 +56,8 @@ contract Integration_Deposit_Delegate_Queue_Complete is IntegrationCheckUtils {
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
 
-        for (uint i = 0; i < withdrawals.length; i++) {
-            uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
+        for (uint256 i = 0; i < withdrawals.length; i++) {
+            uint256[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
             IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
             check_Withdrawal_AsTokens_State(staker, operator, withdrawals[i], strategies, shares, tokens, expectedTokens);
         }
@@ -115,7 +115,7 @@ contract Integration_Deposit_Delegate_Queue_Complete is IntegrationCheckUtils {
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
 
-        for (uint i = 0; i < withdrawals.length; i++) {
+        for (uint256 i = 0; i < withdrawals.length; i++) {
             staker.completeWithdrawalAsShares(withdrawals[i]);
             check_Withdrawal_AsShares_State(staker, operator, withdrawals[i], strategies, shares);
         }
@@ -182,8 +182,8 @@ contract Integration_Deposit_Delegate_Queue_Complete is IntegrationCheckUtils {
         // 4. Complete withdrawals
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
-        for (uint i = 0; i < withdrawals.length; i++) {
-            uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
+        for (uint256 i = 0; i < withdrawals.length; i++) {
+            uint256[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
             IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
             check_Withdrawal_AsTokens_State(staker, operator, withdrawals[i], withdrawStrats, withdrawShares, tokens, expectedTokens);
         }
diff --git a/src/test/integration/tests/Deposit_Delegate_Redelegate_Complete.t.sol b/src/test/integration/tests/Deposit_Delegate_Redelegate_Complete.t.sol
@@ -55,8 +55,10 @@ contract Integration_Deposit_Delegate_Redelegate_Complete is IntegrationCheckUti
         // 4. Complete withdrawal as shares
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
-        staker.completeWithdrawalAsShares(withdrawals[0]);
-        check_Withdrawal_AsShares_Undelegated_State(staker, operator1, withdrawals[0], strategies, shares);
+        for (uint256 i = 0; i < withdrawals.length; ++i) {
+            staker.completeWithdrawalAsShares(withdrawals[i]);
+            check_Withdrawal_AsShares_Undelegated_State(staker, operator1, withdrawals[i], withdrawals[i].strategies, withdrawals[i].shares);
+        }
 
         // 5. Delegate to a new operator
         staker.delegateTo(operator2);
@@ -76,7 +78,7 @@ contract Integration_Deposit_Delegate_Redelegate_Complete is IntegrationCheckUti
         for (uint i = 0; i < withdrawals.length; i++) {
             uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
             IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
-            check_Withdrawal_AsTokens_State(staker, operator2, withdrawals[i], strategies, shares, tokens, expectedTokens);
+            check_Withdrawal_AsTokens_State(staker, operator2, withdrawals[i], withdrawals[i].strategies, withdrawals[i].shares, tokens, expectedTokens);
         }
     }
 
diff --git a/src/test/integration/tests/Deposit_Delegate_Undelegate_Complete.t.sol b/src/test/integration/tests/Deposit_Delegate_Undelegate_Complete.t.sol
@@ -54,9 +54,11 @@ contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationCheckUti
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
 
         // Complete withdrawal
-        uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[0].strategies, withdrawals[0].shares);
-        IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[0]);
-        check_Withdrawal_AsTokens_State(staker, operator, withdrawals[0], strategies, shares, tokens, expectedTokens);
+        for (uint256 i = 0; i < withdrawals.length; ++i) {
+            uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
+            IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
+            check_Withdrawal_AsTokens_State(staker, operator, withdrawals[i], withdrawals[i].strategies, withdrawals[i].shares, tokens, expectedTokens);
+        }
 
         // Check Final State
         assert_HasNoDelegatableShares(staker, "staker should have withdrawn all shares");
@@ -110,8 +112,11 @@ contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationCheckUti
         // 4. Complete withdrawal
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
-        staker.completeWithdrawalAsShares(withdrawals[0]);
-        check_Withdrawal_AsShares_Undelegated_State(staker, operator, withdrawals[0], strategies, shares);
+        for (uint256 i = 0; i < withdrawals.length; ++i) {
+            staker.completeWithdrawalAsShares(withdrawals[i]);
+
+            check_Withdrawal_AsShares_Undelegated_State(staker, operator, withdrawals[i], withdrawals[i].strategies, withdrawals[i].shares);
+        }
 
         // Check final state:
         assert_HasExpectedShares(staker, strategies, shares, "staker should have all original shares");
@@ -161,9 +166,11 @@ contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationCheckUti
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
 
-        uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[0].strategies, withdrawals[0].shares);
-        IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[0]);
-        check_Withdrawal_AsTokens_State(staker, operator, withdrawals[0], strategies, shares, tokens, expectedTokens);
+        for (uint256 i = 0; i < withdrawals.length; ++i) {
+            uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
+            IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
+            check_Withdrawal_AsTokens_State(staker, operator, withdrawals[i], withdrawals[i].strategies, withdrawals[i].shares, tokens, expectedTokens);
+        }
 
         // Check Final State
         assert_HasNoDelegatableShares(staker, "staker should have withdrawn all shares");
@@ -212,8 +219,10 @@ contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationCheckUti
         // 4. Complete withdrawal
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
-        staker.completeWithdrawalAsShares(withdrawals[0]);
-        check_Withdrawal_AsShares_Undelegated_State(staker, operator, withdrawals[0], strategies, shares);
+        for (uint256 i = 0; i < withdrawals.length; ++i) {
+            staker.completeWithdrawalAsShares(withdrawals[i]);
+            check_Withdrawal_AsShares_Undelegated_State(staker, operator, withdrawals[i], withdrawals[i].strategies, withdrawals[i].shares);
+        }
 
         // Check final state:
         assert_HasExpectedShares(staker, strategies, shares, "staker should have all original shares");
diff --git a/src/test/unit/DelegationUnit.t.sol b/src/test/unit/DelegationUnit.t.sol
