Skip to content

Commit

Permalink
test: move assertions into common functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@ypatil12
ypatil12 committed Nov 20, 2023
1 parent 0b062bd commit 7f64588
Show file tree
Hide file tree
Showing 2 changed files with 111 additions and 110 deletions.
141 changes: 33 additions & 108 deletions src/test/integration/tests/Deposit_Delegate_Undelegate_Complete.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,9 @@ 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";

contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationBase {
contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationTestUtils {

/// Randomly generates a user with different held assets. Then:
/// 1. deposit into strategy
Expand Down Expand Up @@ -36,66 +37,28 @@ contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationBase {
assert_HasNoDelegatableShares(staker, "staker should not have delegatable shares before depositing");
assertFalse(delegationManager.isDelegated(address(staker)), "staker should not be delegated");

{
/// 1. 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
staker.depositIntoEigenlayer(strategies, tokenBalances);

assert_HasNoUnderlyingTokenBalance(staker, strategies, "staker should have transferred all underlying tokens");
assert_Snap_AddedStakerShares(staker, strategies, shares, "staker should expected shares in each strategy after depositing");
}

{
/// 2. Delegate to an operator:
//
// ... check that the staker is now delegated to the operator, and that the operator
// was awarded the staker's shares
staker.delegateTo(operator);

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_AddedOperatorShares(operator, strategies, shares, "operator should have received shares");
}
/// 1. Deposit Into Strategies
staker.depositIntoEigenlayer(strategies, tokenBalances);
assertDepositState(staker, strategies, shares);

// 2. Delegate to an operator
staker.delegateTo(operator);
assertDelegationState(staker, operator, strategies, shares);

// 3. Undelegate from an operator
IDelegationManager.Withdrawal memory expectedWithdrawal = _getExpectedWithdrawalStruct(staker);
{
/// 3. 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

bytes32 withdrawalRoot = staker.undelegate();

assertFalse(delegationManager.isDelegated(address(staker)), "staker should not be delegated");
assert_ValidWithdrawalHash(expectedWithdrawal, withdrawalRoot, "calculated withdrawl should match returned root");
assert_withdrawalPending(withdrawalRoot, "staker's withdrawal should now be pending");
assert_Snap_IncrementQueuedWithdrawals(staker, "staker should have increased nonce by 1");
assert_Snap_RemovedOperatorShares(operator, strategies, shares, "failed to remove operator shares");
assert_Snap_RemovedStakerShares(staker, strategies, shares, "failed to remove staker shares");
}
bytes32 withdrawalRoot = staker.undelegate();
assertUndelegateState(staker, operator, expectedWithdrawal, withdrawalRoot, strategies, shares);

// 4. Complete withdrawal
// Fast forward to when we can complete the withdrawal
cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());

{
/// 4. Complete withdrawal(s):
// The staker will complete the withdrawal as tokens
//
// ... check that the withdrawal is not pending, that the staker received the expected tokens, and that the total shares of each
// strategy withdrawn from decreased
uint[] memory expectedTokens = _calculateExpectedTokens(expectedWithdrawal.strategies, expectedWithdrawal.shares);
IERC20[] memory tokens = staker.completeQueuedWithdrawal(expectedWithdrawal, true);

assert_withdrawalNotPending(delegationManager.calculateWithdrawalRoot(expectedWithdrawal), "staker's withdrawal should no longer be pending");
assert_Snap_IncreasedTokenBalances(staker, tokens, expectedTokens, "staker should have received expected tokens");
assert_Snap_DecreasedStrategyShares(strategies, shares, "strategies should have total shares decremented");
}
// Complete withdrawal
uint[] memory expectedTokens = _calculateExpectedTokens(expectedWithdrawal.strategies, expectedWithdrawal.shares);
IERC20[] memory tokens = staker.completeQueuedWithdrawal(expectedWithdrawal, true);

assertWithdrawalAsTokensState(staker, expectedWithdrawal, strategies, shares, tokens, expectedTokens);
}

/// Randomly generates a user with different held assets. Then:
Expand Down Expand Up @@ -128,65 +91,27 @@ contract Integration_Deposit_Delegate_Undelegate_Complete is IntegrationBase {
assert_HasNoDelegatableShares(staker, "staker should not have delegatable shares before depositing");
assertFalse(delegationManager.isDelegated(address(staker)), "staker should not be delegated");

{
/// 1. 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
staker.depositIntoEigenlayer(strategies, tokenBalances);

assert_HasNoUnderlyingTokenBalance(staker, strategies, "staker should have transferred all underlying tokens");
assert_Snap_AddedStakerShares(staker, strategies, shares, "staker should expected shares in each strategy after depositing");
}

{
/// 2. Delegate to an operator:
//
// ... check that the staker is now delegated to the operator, and that the operator
// was awarded the staker's shares
staker.delegateTo(operator);

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_AddedOperatorShares(operator, strategies, shares, "operator should have received shares");
}
/// 1. Deposit Into Strategies
staker.depositIntoEigenlayer(strategies, tokenBalances);
assertDepositState(staker, strategies, shares);

// 2. Delegate to an operator
staker.delegateTo(operator);
assertDelegationState(staker, operator, strategies, shares);

// 3. Undelegate from an operator
IDelegationManager.Withdrawal memory expectedWithdrawal = _getExpectedWithdrawalStruct(staker);
{
/// 3. 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

bytes32 withdrawalRoot = staker.undelegate();

assertFalse(delegationManager.isDelegated(address(staker)), "staker should not be delegated");
assert_ValidWithdrawalHash(expectedWithdrawal, withdrawalRoot, "calculated withdrawl should match returned root");
assert_withdrawalPending(withdrawalRoot, "staker's withdrawal should now be pending");
assert_Snap_IncrementQueuedWithdrawals(staker, "staker should have increased nonce by 1");
assert_Snap_RemovedOperatorShares(operator, strategies, shares, "failed to remove operator shares");
assert_Snap_RemovedStakerShares(staker, strategies, shares, "failed to remove staker shares");
}
bytes32 withdrawalRoot = staker.undelegate();
assertUndelegateState(staker, operator, expectedWithdrawal, withdrawalRoot, strategies, shares);

// 4. Complete withdrawal
// Fast forward to when we can complete the withdrawal
cheats.roll(block.number + delegationManager.withdrawalDelayBlocks());

{
/// 4. Complete withdrawal(s):
// The staker will complete the withdrawal as tokens
//
// ... 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
staker.completeQueuedWithdrawal(expectedWithdrawal, false);

assert_withdrawalNotPending(delegationManager.calculateWithdrawalRoot(expectedWithdrawal), "staker's withdrawal should no longer be pending");
assert_Snap_AddedStakerShares(staker, expectedWithdrawal.strategies, expectedWithdrawal.shares, "staker should have received expected tokens");
assert_Snap_UnchangedStrategyShares(strategies, "strategies should have total shares unchanged");
}
uint[] memory expectedTokens = _calculateExpectedTokens(expectedWithdrawal.strategies, expectedWithdrawal.shares);
IERC20[] memory tokens = staker.completeQueuedWithdrawal(expectedWithdrawal, false);

assertWithdrawalAsSharesState(staker, expectedWithdrawal, strategies, shares);
}

/// @notice Assumes staker and withdrawer are the same
Expand Down
80 changes: 78 additions & 2 deletions src/test/integration/tests/utils.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,83 @@ 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
contract IntegrationTestUtils {
/// @notice Contract that provides utility functions to reuse common test blocks & checks
contract IntegrationTestUtils is IntegrationBase {

function assertDepositState(User staker, IStrategy[] memory strategies, uint256[] 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_AddedStakerShares(staker, strategies, shares, "staker should expected shares in each strategy after depositing");
}

function assertDelegationState(User staker, User operator, IStrategy[] memory strategies, uint256[] 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's 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_AddedOperatorShares(operator, strategies, shares, "operator should have received shares");
}

function assertUndelegateState(
User staker,
User operator,
IDelegationManager.Withdrawal memory withdrawal,
bytes32 withdrawalRoot,
IStrategy[] memory strategies,
uint256[] 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
assertFalse(delegationManager.isDelegated(address(staker)), "staker should not be delegated");
assert_ValidWithdrawalHash(withdrawal, withdrawalRoot, "calculated withdrawl should match returned root");
assert_withdrawalPending(withdrawalRoot, "staker's withdrawal should now be pending");
assert_Snap_IncrementQueuedWithdrawals(staker, "staker should have increased nonce by 1");
assert_Snap_RemovedOperatorShares(operator, strategies, shares, "failed to remove operator shares");
assert_Snap_RemovedStakerShares(staker, strategies, shares, "failed to remove staker shares");
}

function assertWithdrawalAsTokensState(
User staker,
IDelegationManager.Withdrawal memory withdrawal,
IStrategy[] memory strategies,
uint256[] memory shares,
IERC20[] memory tokens,
uint256[] 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 shares, and that the total shares of each
// strategy withdrawn remains unchanged
assert_withdrawalNotPending(delegationManager.calculateWithdrawalRoot(withdrawal), "staker's withdrawal should no longer be pending");
assert_Snap_IncreasedTokenBalances(staker, tokens, expectedTokens, "staker should have received expected tokens");
assert_Snap_DecreasedStrategyShares(strategies, shares, "strategies should have total shares decremented");
}

function assertWithdrawalAsSharesState(
User staker,
IDelegationManager.Withdrawal memory withdrawal,
IStrategy[] memory strategies,
uint256[] 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's withdrawal should no longer be pending");
assert_Snap_AddedStakerShares(staker, strategies, shares, "staker should have received expected shares");
assert_Snap_UnchangedStrategyShares(strategies, "strategies should have total shares unchanged");
}
}

0 comments on commit 7f64588

Please sign in to comment.