Skip to content

Commit

Permalink
test: implement random withdrawal generator (#349)
Browse files Browse the repository at this point in the history 
see PR for changes/notes
  • Loading branch information
@wadealexc @ypatil12
wadealexc authored and ypatil12 committed Dec 4, 2023
1 parent 102aceb commit 64eafb3
Show file tree
Hide file tree
Showing 4 changed files with 540 additions and 67 deletions.
224 changes: 199 additions & 25 deletions src/test/integration/IntegrationBase.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,8 @@ abstract contract IntegrationBase is IntegrationDeployer {
operator.registerAsOperator();
operator.depositIntoEigenlayer(strategies, tokenBalances);

assert_Snap_AddedStakerShares(operator, strategies, tokenBalances, "_newRandomOperator: failed to add delegatable shares");
assert_Snap_AddedOperatorShares(operator, strategies, tokenBalances, "_newRandomOperator: failed to award shares to operator");
assert_Snap_Added_StakerShares(operator, strategies, tokenBalances, "_newRandomOperator: failed to add delegatable shares");
assert_Snap_Added_OperatorShares(operator, strategies, tokenBalances, "_newRandomOperator: failed to award shares to operator");
assertTrue(delegationManager.isOperator(address(operator)), "_newRandomOperator: operator should be registered");

return (operator, strategies, tokenBalances);
Expand Down Expand Up @@ -120,12 +120,14 @@ 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++) {
assert_withdrawalPending(withdrawalRoots[i], err);
}
}

<<<<<<< HEAD
function assert_withdrawalPending(bytes32 withdrawalRoot, string memory err) internal {
assertTrue(delegationManager.pendingWithdrawals(withdrawalRoot), err);
}
Expand All @@ -140,13 +142,29 @@ abstract contract IntegrationBase is IntegrationDeployer {
assertFalse(delegationManager.pendingWithdrawals(withdrawalRoot), 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);
}
}

/// @dev Asserts that the hash of each withdrawal corresponds to the provided withdrawal root
>>>>>>> 272a7e2f (test: implement random withdrawal generator (#349))
function assert_ValidWithdrawalHashes(
IDelegationManager.Withdrawal[] memory withdrawals,
bytes32[] memory withdrawalRoots,
string memory err
) internal {
bytes32[] memory expectedRoots = _getWithdrawalHashes(withdrawals);

for (uint i = 0; i < withdrawals.length; i++) {
<<<<<<< HEAD
assert_ValidWithdrawalHash(withdrawals[i], withdrawalRoots[i], err);
=======
assertEq(withdrawalRoots[i], expectedRoots[i], err);
>>>>>>> 272a7e2f (test: implement random withdrawal generator (#349))
}
}

Expand All @@ -158,14 +176,16 @@ abstract contract IntegrationBase is IntegrationDeployer {
assertEq(withdrawalRoot, delegationManager.calculateWithdrawalRoot(withdrawal), err);
}

/**
* Snapshot assertions combine Timemachine's snapshots with assertions
* that allow easy comparisons between prev/cur values
*/
/*******************************************************************************
SNAPSHOT ASSERTIONS
TIME TRAVELERS ONLY BEYOND THIS POINT
*******************************************************************************/

/// @dev Check that the operator has `addedShares` additional shares for each
/// strategy since the last snapshot
function assert_Snap_AddedOperatorShares(
/// Snapshot assertions for delegationManager.operatorShares:

/// @dev Check that the operator has `addedShares` additional operator shares
// for each strategy since the last snapshot
function assert_Snap_Added_OperatorShares(
User operator,
IStrategy[] memory strategies,
uint[] memory addedShares,
Expand All @@ -181,9 +201,9 @@ abstract contract IntegrationBase is IntegrationDeployer {
}
}

/// @dev Check that the operator has `removedShares` prior shares for each
/// strategy since the last snapshot
function assert_Snap_RemovedOperatorShares(
/// @dev Check that the operator has `removedShares` fewer operator shares
/// for each strategy since the last snapshot
function assert_Snap_Removed_OperatorShares(
User operator,
IStrategy[] memory strategies,
uint[] memory removedShares,
Expand All @@ -199,9 +219,29 @@ abstract contract IntegrationBase is IntegrationDeployer {
}
}

/// @dev Check that the staker has `addedShares` additional shares for each
/// strategy since the last snapshot
function assert_Snap_AddedStakerShares(
/// @dev Check that the operator's shares in ALL strategies have not changed
/// since the last snapshot
function assert_Snap_Unchanged_OperatorShares(
User operator,
string memory err
) internal {
IStrategy[] memory strategies = allStrats;

uint[] memory curShares = _getOperatorShares(operator, strategies);
// Use timewarp to get previous operator shares
uint[] memory prevShares = _getPrevOperatorShares(operator, strategies);

// For each strategy, check (prev == cur)
for (uint i = 0; i < strategies.length; i++) {
assertEq(prevShares[i], curShares[i], err);
}
}

/// Snapshot assertions for strategyMgr.stakerStrategyShares and eigenPodMgr.podOwnerShares:

/// @dev Check that the staker has `addedShares` additional delegatable shares
/// for each strategy since the last snapshot
function assert_Snap_Added_StakerShares(
User staker,
IStrategy[] memory strategies,
uint[] memory addedShares,
Expand All @@ -217,9 +257,9 @@ abstract contract IntegrationBase is IntegrationDeployer {
}
}

/// @dev Check that the staker has `removedShares` prior shares for each
/// strategy since the last snapshot
function assert_Snap_RemovedStakerShares(
/// @dev Check that the staker has `removedShares` fewer delegatable shares
/// for each strategy since the last snapshot
function assert_Snap_Removed_StakerShares(
User staker,
IStrategy[] memory strategies,
uint[] memory removedShares,
Expand All @@ -235,18 +275,25 @@ abstract contract IntegrationBase is IntegrationDeployer {
}
}

function assert_Snap_IncreasedQueuedWithdrawals(
User staker,
IDelegationManager.Withdrawal[] memory withdrawals,
/// @dev Check that the staker's delegatable shares in ALL strategies have not changed
/// since the last snapshot
function assert_Snap_Unchanged_StakerShares(
User staker,
string memory err
) internal {
uint curQueuedWithdrawals = _getCumulativeWithdrawals(staker);
// Use timewarp to get previous cumulative withdrawals
uint prevQueuedWithdrawals = _getPrevCumulativeWithdrawals(staker);
IStrategy[] memory strategies = allStrats;

assertEq(prevQueuedWithdrawals + withdrawals.length, curQueuedWithdrawals, err);
uint[] memory curShares = _getStakerShares(staker, strategies);
// Use timewarp to get previous staker shares
uint[] memory prevShares = _getPrevStakerShares(staker, strategies);

// For each strategy, check (prev == cur)
for (uint i = 0; i < strategies.length; i++) {
assertEq(prevShares[i], curShares[i], err);
}
}

<<<<<<< HEAD
function assert_Snap_IncrementQueuedWithdrawals(
User staker,
string memory err
Expand All @@ -259,6 +306,13 @@ abstract contract IntegrationBase is IntegrationDeployer {
}

function assert_Snap_IncreasedTokenBalances(
=======
/// Snapshot assertions for underlying token balances:

/// @dev Check that the staker has `addedTokens` additional underlying tokens
// since the last snapshot
function assert_Snap_Added_TokenBalances(
>>>>>>> 272a7e2f (test: implement random withdrawal generator (#349))
User staker,
IERC20[] memory tokens,
uint[] memory addedTokens,
Expand All @@ -276,6 +330,7 @@ abstract contract IntegrationBase is IntegrationDeployer {
}
}

<<<<<<< HEAD
function assert_Snap_DecreasedStrategyShares(
IStrategy[] memory strategies,
uint256[] memory removedShares,
Expand Down Expand Up @@ -318,6 +373,96 @@ abstract contract IntegrationBase is IntegrationDeployer {
/**
* Helpful getters:
*/
=======
/// @dev Check that the staker has `removedTokens` fewer underlying tokens
// since the last snapshot
function assert_Snap_Removed_TokenBalances(
User staker,
IStrategy[] memory strategies,
uint[] memory removedTokens,
string memory err
) internal {
IERC20[] memory tokens = _getUnderlyingTokens(strategies);

uint[] memory curTokenBalances = _getTokenBalances(staker, tokens);
// Use timewarp to get previous token balances
uint[] memory prevTokenBalances = _getPrevTokenBalances(staker, tokens);

for (uint i = 0; i < tokens.length; i++) {
uint prevBalance = prevTokenBalances[i];
uint curBalance = curTokenBalances[i];

assertEq(prevBalance - removedTokens[i], curBalance, err);
}
}

/// @dev Check that the staker's underlying token balance for ALL tokens have
/// not changed since the last snapshot
function assert_Snap_Unchanged_TokenBalances(
User staker,
string memory err
) internal {
IERC20[] memory tokens = allTokens;

uint[] memory curTokenBalances = _getTokenBalances(staker, tokens);
// Use timewarp to get previous token balances
uint[] memory prevTokenBalances = _getPrevTokenBalances(staker, tokens);

for (uint i = 0; i < tokens.length; i++) {
assertEq(prevTokenBalances[i], curTokenBalances[i], err);
}
}

/// Other snapshot assertions:

function assert_Snap_Added_QueuedWithdrawals(
User staker,
IDelegationManager.Withdrawal[] memory withdrawals,
string memory err
) internal {
uint curQueuedWithdrawals = _getCumulativeWithdrawals(staker);
// Use timewarp to get previous cumulative withdrawals
uint prevQueuedWithdrawals = _getPrevCumulativeWithdrawals(staker);

assertEq(prevQueuedWithdrawals + withdrawals.length, curQueuedWithdrawals, err);
}

/*******************************************************************************
UTILITY METHODS
*******************************************************************************/

function _randWithdrawal(
IStrategy[] memory strategies,
uint[] memory shares
) internal returns (IStrategy[] memory, uint[] memory) {
uint stratsToWithdraw = _randUint({ min: 1, max: strategies.length });

IStrategy[] memory withdrawStrats = new IStrategy[](stratsToWithdraw);
uint[] memory withdrawShares = new uint[](stratsToWithdraw);

for (uint i = 0; i < stratsToWithdraw; i++) {
uint sharesToWithdraw;

if (strategies[i] == BEACONCHAIN_ETH_STRAT) {
// For native eth, withdraw a random amount of gwei (at least 1)
uint portion = _randUint({ min: 1, max: shares[i] / GWEI_TO_WEI });
portion *= GWEI_TO_WEI;

sharesToWithdraw = shares[i] - portion;
} else {
// For LSTs, withdraw a random amount of shares (at least 1)
uint portion = _randUint({ min: 1, max: shares[i] });

sharesToWithdraw = shares[i] - portion;
}

withdrawStrats[i] = strategies[i];
withdrawShares[i] = sharesToWithdraw;
}

return (withdrawStrats, withdrawShares);
}
>>>>>>> 272a7e2f (test: implement random withdrawal generator (#349))

/// @dev For some strategies/underlying token balances, calculate the expected shares received
/// from depositing all tokens
Expand Down Expand Up @@ -356,6 +501,35 @@ abstract contract IntegrationBase is IntegrationDeployer {
return expectedTokens;
}

function _getWithdrawalHashes(
IDelegationManager.Withdrawal[] memory withdrawals
) internal view returns (bytes32[] memory) {
bytes32[] memory withdrawalRoots = new bytes32[](withdrawals.length);

for (uint i = 0; i < withdrawals.length; i++) {
withdrawalRoots[i] = delegationManager.calculateWithdrawalRoot(withdrawals[i]);
}

return withdrawalRoots;
}

/// @dev Converts a list of strategies to underlying tokens
function _getUnderlyingTokens(IStrategy[] memory strategies) internal view returns (IERC20[] memory) {
IERC20[] memory tokens = new IERC20[](strategies.length);

for (uint i = 0; i < tokens.length; i++) {
IStrategy strat = strategies[i];

if (strat == BEACONCHAIN_ETH_STRAT) {
tokens[i] = NATIVE_ETH;
} else {
tokens[i] = strat.underlyingToken();
}
}

return tokens;
}

modifier timewarp() {
uint curState = timeMachine.warpToLast();
_;
Expand Down
12 changes: 8 additions & 4 deletions src/test/integration/IntegrationDeployer.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,8 @@ abstract contract IntegrationDeployer is Test, IUserDeployer {
// which of these lists to select user assets from.
IStrategy[] lstStrats;
IStrategy[] ethStrats; // only has one strat tbh
IStrategy[] mixedStrats; // just a combination of the above 2 lists
IStrategy[] allStrats; // just a combination of the above 2 lists
IERC20[] allTokens; // `allStrats`, but contains all of the underlying tokens instead

// Mock Contracts to deploy
ETHPOSDepositMock ethPOSDeposit;
Expand Down Expand Up @@ -80,6 +81,7 @@ abstract contract IntegrationDeployer is Test, IUserDeployer {

uint constant MIN_BALANCE = 1e6;
uint constant MAX_BALANCE = 5e6;
uint constant GWEI_TO_WEI = 1e9;

// Flags
uint constant FLAG = 1;
Expand Down Expand Up @@ -256,7 +258,8 @@ abstract contract IntegrationDeployer is Test, IUserDeployer {
_newStrategyAndToken("Strategy3Token", "str3", 10e50, address(this)); // initialSupply, owner

ethStrats.push(BEACONCHAIN_ETH_STRAT);
mixedStrats.push(BEACONCHAIN_ETH_STRAT);
allStrats.push(BEACONCHAIN_ETH_STRAT);
allTokens.push(NATIVE_ETH);

// Create time machine and set block timestamp forward so we can create EigenPod proofs in the past
timeMachine = new TimeMachine();
Expand Down Expand Up @@ -286,9 +289,10 @@ abstract contract IntegrationDeployer is Test, IUserDeployer {
cheats.prank(strategyManager.strategyWhitelister());
strategyManager.addStrategiesToDepositWhitelist(strategies);

// Add to lstStrats and mixedStrats
// Add to lstStrats and allStrats
lstStrats.push(strategy);
mixedStrats.push(strategy);
allStrats.push(strategy);
allTokens.push(underlyingToken);
}

function _configRand(
Expand Down
8 changes: 3 additions & 5 deletions src/test/integration/User.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -137,13 +137,11 @@ contract User is Test {
function queueWithdrawals(
IStrategy[] memory strategies,
uint[] memory shares
) public createSnapshot virtual returns (IDelegationManager.Withdrawal[] memory, bytes32[] memory) {
) public createSnapshot virtual returns (IDelegationManager.Withdrawal[] memory) {

address operator = delegationManager.delegatedTo(address(this));
address withdrawer = address(this);
uint nonce = delegationManager.cumulativeWithdrawalsQueued(address(this));

bytes32[] memory withdrawalRoots;

// Create queueWithdrawals params
IDelegationManager.QueuedWithdrawalParams[] memory params = new IDelegationManager.QueuedWithdrawalParams[](1);
Expand All @@ -165,12 +163,12 @@ contract User is Test {
shares: shares
});

withdrawalRoots = delegationManager.queueWithdrawals(params);
bytes32[] memory withdrawalRoots = delegationManager.queueWithdrawals(params);

// Basic sanity check - we do all other checks outside this file
assertEq(withdrawals.length, withdrawalRoots.length, "User.queueWithdrawals: length mismatch");

return (withdrawals, withdrawalRoots);
return (withdrawals);
}

function completeQueuedWithdrawal(
Expand Down
Loading

0 comments on commit 64eafb3

Please sign in to comment.