test: add separate withdrawal as tokens and shares methods

Author: ypatil12

src/test/integration/tests/utils.t.sol renamed to src/test/integration/IntegrationChecks.t.sol

@@ -5,7 +5,7 @@ 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 IntegrationTestUtils is IntegrationBase {
+contract IntegrationCheckUtils is IntegrationBase {
 
     function check_Deposit_State(User staker, IStrategy[] memory strategies, uint[] memory shares) internal {
         /// Deposit into strategies:

src/test/integration/User.t.sol

@@ -176,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++) {

src/test/integration/tests/Deposit_Delegate_Queue_Complete.t.sol

@@ -1,10 +1,10 @@
 // SPDX-License-Identifier: BUSL-1.1
 pragma solidity =0.8.12;
 
-import "src/test/integration/tests/utils.t.sol";
+import "src/test/integration/IntegrationChecks.t.sol";
 import "src/test/integration/User.t.sol";
 
-contract Integration_Deposit_Delegate_Queue_Complete is IntegrationTestUtils {
+contract Integration_Deposit_Delegate_Queue_Complete is IntegrationCheckUtils {
 
     /*******************************************************************************
                                 FULL WITHDRAWALS
@@ -58,7 +58,7 @@ contract Integration_Deposit_Delegate_Queue_Complete is IntegrationTestUtils {
 
         for (uint i = 0; i < withdrawals.length; i++) {
             uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
-            IERC20[] memory tokens = staker.completeQueuedWithdrawal(withdrawals[i], true);
+            IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
             check_Withdrawal_AsTokens_State(staker, operator, withdrawals[i], strategies, shares, tokens, expectedTokens);
         }
 
@@ -116,7 +116,7 @@ contract Integration_Deposit_Delegate_Queue_Complete is IntegrationTestUtils {
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
 
         for (uint i = 0; i < withdrawals.length; i++) {
-            staker.completeQueuedWithdrawal(withdrawals[i], false);
+            staker.completeWithdrawalAsShares(withdrawals[i]);
             check_Withdrawal_AsShares_State(staker, operator, withdrawals[i], strategies, shares);
         }
 
@@ -184,7 +184,7 @@ contract Integration_Deposit_Delegate_Queue_Complete is IntegrationTestUtils {
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
         for (uint i = 0; i < withdrawals.length; i++) {
             uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
-            IERC20[] memory tokens = staker.completeQueuedWithdrawal(withdrawals[i], true);
+            IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
             check_Withdrawal_AsTokens_State(staker, operator, withdrawals[i], withdrawStrats, withdrawShares, tokens, expectedTokens);
         }
 
@@ -246,7 +246,7 @@ contract Integration_Deposit_Delegate_Queue_Complete is IntegrationTestUtils {
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
 
         for (uint i = 0; i < withdrawals.length; i++) {
-            staker.completeQueuedWithdrawal(withdrawals[i], false);
+            staker.completeWithdrawalAsShares(withdrawals[i]);
             check_Withdrawal_AsShares_State(staker, operator, withdrawals[i], withdrawStrats, withdrawShares);
         }
 

src/test/integration/tests/Deposit_Delegate_Redelegate_Complete.t.sol

@@ -1,11 +1,10 @@
 // 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";
-import "src/test/integration/tests/utils.t.sol";
+import "src/test/integration/IntegrationChecks.t.sol";
 
-contract Integration_Deposit_Delegate_Redelegate_Complete is IntegrationTestUtils {
+contract Integration_Deposit_Delegate_Redelegate_Complete is IntegrationCheckUtils {
     /// Randomly generates a user with different held assets. Then:
     /// 1. deposit into strategy
     /// 2. delegate to an operator
@@ -56,7 +55,7 @@ contract Integration_Deposit_Delegate_Redelegate_Complete is IntegrationTestUtil
         // 4. Complete withdrawal as shares
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
-        staker.completeQueuedWithdrawal(withdrawals[0], false);
+        staker.completeWithdrawalAsShares(withdrawals[0]);
         check_Withdrawal_AsShares_Undelegated_State(staker, operator1, withdrawals[0], strategies, shares);
 
         // 5. Delegate to a new operator
@@ -76,7 +75,7 @@ contract Integration_Deposit_Delegate_Redelegate_Complete is IntegrationTestUtil
         // Complete withdrawals
         for (uint i = 0; i < withdrawals.length; i++) {
             uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[i].strategies, withdrawals[i].shares);
-            IERC20[] memory tokens = staker.completeQueuedWithdrawal(withdrawals[i], true);
+            IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[i]);
             check_Withdrawal_AsTokens_State(staker, operator2, withdrawals[i], strategies, shares, tokens, expectedTokens);
         }
     }

src/test/integration/tests/Deposit_Delegate_Undelegate_Complete.t.sol

@@ -2,9 +2,9 @@
 pragma solidity =0.8.12;
 
 import "src/test/integration/User.t.sol";
-import "src/test/integration/tests/utils.t.sol";
+import "src/test/integration/IntegrationChecks.t.sol";
 
-contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationTestUtils {
+contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationCheckUtils {
 
     /// Randomly generates a user with different held assets. Then:
     /// 1. deposit into strategy
@@ -55,7 +55,7 @@ contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationTestUtil
 
         // Complete withdrawal
         uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[0].strategies, withdrawals[0].shares);
-        IERC20[] memory tokens = staker.completeQueuedWithdrawal(withdrawals[0], true);
+        IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[0]);
         check_Withdrawal_AsTokens_State(staker, operator, withdrawals[0], strategies, shares, tokens, expectedTokens);
 
         // Check Final State
@@ -110,7 +110,7 @@ contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationTestUtil
         // 4. Complete withdrawal
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
-        staker.completeQueuedWithdrawal(withdrawals[0], false);
+        staker.completeWithdrawalAsShares(withdrawals[0]);
         check_Withdrawal_AsShares_Undelegated_State(staker, operator, withdrawals[0], strategies, shares);
 
         // Check final state:
@@ -162,7 +162,7 @@ contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationTestUtil
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
 
         uint[] memory expectedTokens = _calculateExpectedTokens(withdrawals[0].strategies, withdrawals[0].shares);
-        IERC20[] memory tokens = staker.completeQueuedWithdrawal(withdrawals[0], true);
+        IERC20[] memory tokens = staker.completeWithdrawalAsTokens(withdrawals[0]);
         check_Withdrawal_AsTokens_State(staker, operator, withdrawals[0], strategies, shares, tokens, expectedTokens);
 
         // Check Final State
@@ -212,7 +212,7 @@ contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationTestUtil
         // 4. Complete withdrawal
         // Fast forward to when we can complete the withdrawal
         cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());
-        staker.completeQueuedWithdrawal(withdrawals[0], false);
+        staker.completeWithdrawalAsShares(withdrawals[0]);
         check_Withdrawal_AsShares_Undelegated_State(staker, operator, withdrawals[0], strategies, shares);
 
         // Check final state: