Skip to content

Commit

Permalink
test: support random balance updates in integration tests (#364)
Browse files Browse the repository at this point in the history 
- also adds user-level logging
  • Loading branch information
@wadealexc
wadealexc authored Dec 5, 2023
1 parent c439468 commit e57b40a
Show file tree
Hide file tree
Showing 6 changed files with 464 additions and 19 deletions.
190 changes: 181 additions & 9 deletions src/test/integration/IntegrationBase.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,19 @@ pragma solidity =0.8.12;
import "forge-std/Test.sol";

import "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";
import "@openzeppelin/contracts/utils/Strings.sol";

import "src/test/integration/IntegrationDeployer.t.sol";
import "src/test/integration/TimeMachine.t.sol";
import "src/test/integration/User.t.sol";

abstract contract IntegrationBase is IntegrationDeployer {

using Strings for *;

uint numStakers = 0;
uint numOperators = 0;

/**
* Gen/Init methods:
*/
Expand All @@ -20,15 +26,21 @@ abstract contract IntegrationBase is IntegrationDeployer {
* This user is ready to deposit into some strategies and has some underlying token balances
*/
function _newRandomStaker() internal returns (User, IStrategy[] memory, uint[] memory) {
(User staker, IStrategy[] memory strategies, uint[] memory tokenBalances) = _randUser();
string memory stakerName = string.concat("- Staker", numStakers.toString());
numStakers++;

(User staker, IStrategy[] memory strategies, uint[] memory tokenBalances) = _randUser(stakerName);

assert_HasUnderlyingTokenBalances(staker, strategies, tokenBalances, "_newRandomStaker: failed to award token balances");

return (staker, strategies, tokenBalances);
}

function _newRandomOperator() internal returns (User, IStrategy[] memory, uint[] memory) {
(User operator, IStrategy[] memory strategies, uint[] memory tokenBalances) = _randUser();
string memory operatorName = string.concat("- Operator", numOperators.toString());
numOperators++;

(User operator, IStrategy[] memory strategies, uint[] memory tokenBalances) = _randUser(operatorName);

operator.registerAsOperator();
operator.depositIntoEigenlayer(strategies, tokenBalances);
Expand Down Expand Up @@ -148,8 +160,6 @@ abstract contract IntegrationBase is IntegrationDeployer {
bytes32[] memory withdrawalRoots,
string memory err
) internal {
bytes32[] memory expectedRoots = _getWithdrawalHashes(withdrawals);

for (uint i = 0; i < withdrawals.length; i++) {
assert_ValidWithdrawalHash(withdrawals[i], withdrawalRoots[i], err);
}
Expand Down Expand Up @@ -224,6 +234,28 @@ abstract contract IntegrationBase is IntegrationDeployer {
}
}

function assert_Snap_Delta_OperatorShares(
User operator,
IStrategy[] memory strategies,
int[] memory shareDeltas,
string memory err
) internal {
uint[] memory curShares = _getOperatorShares(operator, strategies);
// Use timewarp to get previous operator shares
uint[] memory prevShares = _getPrevOperatorShares(operator, strategies);

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

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

/// @dev Check that the staker has `addedShares` additional delegatable shares
Expand Down Expand Up @@ -319,6 +351,22 @@ abstract contract IntegrationBase is IntegrationDeployer {
}
}

function assert_Snap_Delta_StakerShares(
User staker,
IStrategy[] memory strategies,
int[] memory shareDeltas,
string memory err
) internal {
int[] memory curShares = _getStakerSharesInt(staker, strategies);
// Use timewarp to get previous staker shares
int[] memory prevShares = _getPrevStakerSharesInt(staker, strategies);

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

/// Snapshot assertions for underlying token balances:

/// @dev Check that the staker has `addedTokens` additional underlying tokens
Expand Down Expand Up @@ -396,7 +444,6 @@ abstract contract IntegrationBase is IntegrationDeployer {

function assert_Snap_Added_QueuedWithdrawal(
User staker,
IDelegationManager.Withdrawal memory withdrawal,
string memory err
) internal {
uint curQueuedWithdrawal = _getCumulativeWithdrawals(staker);
Expand Down Expand Up @@ -442,9 +489,109 @@ abstract contract IntegrationBase is IntegrationDeployer {
return (withdrawStrats, withdrawShares);
}

/**
* Helpful getters:
*/
function _randBalanceUpdate(
User staker,
IStrategy[] memory strategies
) internal returns (int[] memory, int[] memory, int[] memory) {

int[] memory tokenDeltas = new int[](strategies.length);
int[] memory stakerShareDeltas = new int[](strategies.length);
int[] memory operatorShareDeltas = new int[](strategies.length);

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

if (strat == BEACONCHAIN_ETH_STRAT) {
// TODO - could choose and set a "next updatable validator" at random here
uint40 validator = staker.getUpdatableValidator();
uint64 beaconBalanceGwei = beaconChain.balanceOfGwei(validator);

// For native eth, add or remove a random amount of Gwei - minimum 1
// and max of the current beacon chain balance
int64 deltaGwei = int64(int(_randUint({ min: 1, max: beaconBalanceGwei })));
bool addTokens = _randBool();
deltaGwei = addTokens ? deltaGwei : -deltaGwei;

tokenDeltas[i] = int(deltaGwei) * int(GWEI_TO_WEI);

// stakerShareDeltas[i] = _calculateSharesDelta(newPodBalanceGwei, oldPodBalanceGwei);
stakerShareDeltas[i] = _calcNativeETHStakerShareDelta(staker, validator, beaconBalanceGwei, deltaGwei);
operatorShareDeltas[i] = _calcNativeETHOperatorShareDelta(staker, stakerShareDeltas[i]);

emit log_named_uint("current beacon balance (gwei): ", beaconBalanceGwei);
// emit log_named_uint("current validator pod balance (gwei): ", oldPodBalanceGwei);
emit log_named_int("beacon balance delta (gwei): ", deltaGwei);
emit log_named_int("staker share delta (gwei): ", stakerShareDeltas[i] / int(GWEI_TO_WEI));
emit log_named_int("operator share delta (gwei): ", operatorShareDeltas[i] / int(GWEI_TO_WEI));
} else {
// For LSTs, mint a random token amount
uint portion = _randUint({ min: MIN_BALANCE, max: MAX_BALANCE });
StdCheats.deal(address(strat.underlyingToken()), address(staker), portion);

int delta = int(portion);
tokenDeltas[i] = delta;
stakerShareDeltas[i] = int(strat.underlyingToShares(uint(delta)));
operatorShareDeltas[i] = int(strat.underlyingToShares(uint(delta)));
}
}
return (tokenDeltas, stakerShareDeltas, operatorShareDeltas);
}

function _calcNativeETHStakerShareDelta(
User staker,
uint40 validatorIndex,
uint64 beaconBalanceGwei,
int64 deltaGwei
) internal view returns (int) {
uint64 oldPodBalanceGwei =
staker
.pod()
.validatorPubkeyHashToInfo(beaconChain.pubkeyHash(validatorIndex))
.restakedBalanceGwei;

uint64 newPodBalanceGwei = _calcPodBalance(beaconBalanceGwei, deltaGwei);

return (int(uint(newPodBalanceGwei)) - int(uint(oldPodBalanceGwei))) * int(GWEI_TO_WEI);
}

function _calcPodBalance(uint64 beaconBalanceGwei, int64 deltaGwei) internal pure returns (uint64) {
uint64 podBalanceGwei;
if (deltaGwei < 0) {
podBalanceGwei = beaconBalanceGwei - uint64(uint(int(-deltaGwei)));
} else {
podBalanceGwei = beaconBalanceGwei + uint64(uint(int(deltaGwei)));
}

if (podBalanceGwei > MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR) {
podBalanceGwei = MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR;
}

return podBalanceGwei;
}

function _calcNativeETHOperatorShareDelta(User staker, int shareDelta) internal view returns (int) {
int curPodOwnerShares = eigenPodManager.podOwnerShares(address(staker));
int newPodOwnerShares = curPodOwnerShares + shareDelta;

if (curPodOwnerShares <= 0) {
// if the shares started negative and stayed negative, then there cannot have been an increase in delegateable shares
if (newPodOwnerShares <= 0) {
return 0;
// if the shares started negative and became positive, then the increase in delegateable shares is the ending share amount
} else {
return newPodOwnerShares;
}
} else {
// if the shares started positive and became negative, then the decrease in delegateable shares is the starting share amount
if (newPodOwnerShares <= 0) {
return (-curPodOwnerShares);
// if the shares started positive and stayed positive, then the change in delegateable shares
// is the difference between starting and ending amounts
} else {
return (newPodOwnerShares - curPodOwnerShares);
}
}
}

/// @dev For some strategies/underlying token balances, calculate the expected shares received
/// from depositing all tokens
Expand All @@ -456,7 +603,7 @@ abstract contract IntegrationBase is IntegrationDeployer {

uint tokenBalance = tokenBalances[i];
if (strat == BEACONCHAIN_ETH_STRAT) {
expectedShares[i] = tokenBalances[i];
expectedShares[i] = tokenBalance;
} else {
expectedShares[i] = strat.underlyingToShares(tokenBalance);
}
Expand Down Expand Up @@ -570,6 +717,31 @@ abstract contract IntegrationBase is IntegrationDeployer {
return curShares;
}

/// @dev Uses timewarp modifier to get staker shares at the last snapshot
function _getPrevStakerSharesInt(
User staker,
IStrategy[] memory strategies
) internal timewarp() returns (int[] memory) {
return _getStakerSharesInt(staker, strategies);
}

/// @dev Looks up each strategy and returns a list of the staker's shares
function _getStakerSharesInt(User staker, IStrategy[] memory strategies) internal view returns (int[] memory) {
int[] memory curShares = new int[](strategies.length);

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

if (strat == BEACONCHAIN_ETH_STRAT) {
curShares[i] = eigenPodManager.podOwnerShares(address(staker));
} else {
curShares[i] = int(strategyManager.stakerStrategyShares(address(staker), strat));
}
}

return curShares;
}

function _getPrevCumulativeWithdrawals(User staker) internal timewarp() returns (uint) {
return _getCumulativeWithdrawals(staker);
}
Expand Down
23 changes: 20 additions & 3 deletions src/test/integration/IntegrationChecks.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,11 @@ import "src/test/integration/User.t.sol";
/// @notice Contract that provides utility functions to reuse common test blocks & checks
contract IntegrationCheckUtils is IntegrationBase {

function check_Deposit_State(User staker, IStrategy[] memory strategies, uint[] memory shares) internal {
function check_Deposit_State(
User staker,
IStrategy[] memory strategies,
uint[] 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.
Expand All @@ -18,18 +22,31 @@ contract IntegrationCheckUtils is IntegrationBase {
assert_Snap_Added_StakerShares(staker, strategies, shares, "staker should expected shares in each strategy after depositing");
}

function check_Delegation_State(User staker, User operator, IStrategy[] memory strategies, uint[] memory shares) internal {
function check_Delegation_State(
User staker,
User operator,
IStrategy[] memory strategies,
uint[] 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 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_Unchanged_StakerShares(staker, "staker shares should be unchanged after delegating");
assert_Snap_Added_OperatorShares(operator, strategies, shares, "operator should have received shares");
}

function check_QueuedWithdrawal_State(User staker, User operator, IStrategy[] memory strategies, uint[] memory shares, IDelegationManager.Withdrawal[] memory withdrawals, bytes32[] memory withdrawalRoots) internal {
function check_QueuedWithdrawal_State(
User staker,
User operator,
IStrategy[] memory strategies,
uint[] memory shares,
IDelegationManager.Withdrawal[] memory withdrawals,
bytes32[] memory withdrawalRoots
) internal {
// The staker will queue one or more withdrawals for the selected strategies and shares
//
// ... check that each withdrawal was successfully enqueued, that the returned roots
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 @@ -170,7 +170,7 @@ abstract contract IntegrationDeployer is Test, IUserDeployer {
delayedWithdrawalRouter,
eigenPodManager,
MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR,
GOERLI_GENESIS_TIME
0
);

eigenPodBeacon = new UpgradeableBeacon(address(pod));
Expand Down Expand Up @@ -329,7 +329,7 @@ abstract contract IntegrationDeployer is Test, IUserDeployer {
*
* Assets are pulled from `strategies` based on a random staker/operator `assetType`
*/
function _randUser() internal returns (User, IStrategy[] memory, uint[] memory) {
function _randUser(string memory name) internal returns (User, IStrategy[] memory, uint[] memory) {
// For the new user, select what type of assets they'll have and whether
// they'll use `xWithSignature` methods.
//
Expand All @@ -340,11 +340,11 @@ abstract contract IntegrationDeployer is Test, IUserDeployer {
// Create User contract based on deposit type:
User user;
if (userType == DEFAULT) {
user = new User();
user = new User(name);
} else if (userType == ALT_METHODS) {
// User will use nonstandard methods like:
// `delegateToBySignature` and `depositIntoStrategyWithSignature`
user = User(new User_AltMethods());
user = User(new User_AltMethods(name));
} else {
revert("_randUser: unimplemented userType");
}
Expand Down Expand Up @@ -462,6 +462,10 @@ abstract contract IntegrationDeployer is Test, IUserDeployer {
return min + value;
}

function _randBool() internal returns (bool) {
return _randUint({ min: 0, max: 1 }) == 0;
}

function _randAssetType() internal returns (uint) {
uint idx = _randUint({ min: 0, max: assetTypes.length - 1 });
uint assetType = uint(uint8(assetTypes[idx]));
Expand Down
Loading

0 comments on commit e57b40a

Please sign in to comment.