Skip to content

Commit 6ba8d82

Browse files
wadealexcypatil12
authored andcommitted
test: enable shared setups for integration tests (#1036)
* test: improve integration invariants * also removes unneeded fork logic * adds checks to some invariants * fixes some broken tests * test(integration): enable shared setups
1 parent 2da77ec commit 6ba8d82

9 files changed

+201
-319
lines changed

src/test/integration/IntegrationBase.t.sol

Lines changed: 74 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ abstract contract IntegrationBase is IntegrationDeployer {
2020
using Strings for *;
2121
using print for *;
2222

23-
using ArrayLib for IStrategy[];
23+
using ArrayLib for *;
2424

2525
uint numStakers = 0;
2626
uint numOperators = 0;
@@ -49,7 +49,7 @@ abstract contract IntegrationBase is IntegrationDeployer {
4949
IStrategy[] memory strategies;
5050
uint[] memory tokenBalances;
5151

52-
if (forkType == MAINNET && !isUpgraded) {
52+
if (!isUpgraded) {
5353
stakerName = string.concat("M2Staker", cheats.toString(numStakers));
5454

5555
(staker, strategies, tokenBalances) = _randUser(stakerName);
@@ -79,7 +79,7 @@ abstract contract IntegrationBase is IntegrationDeployer {
7979
uint[] memory tokenBalances;
8080
uint[] memory addedShares;
8181

82-
if (forkType == MAINNET && !isUpgraded) {
82+
if (!isUpgraded) {
8383
string memory operatorName = string.concat("M2Operator", numOperators.toString());
8484

8585
// Create an operator for M2.
@@ -102,11 +102,12 @@ abstract contract IntegrationBase is IntegrationDeployer {
102102
operator.registerAsOperator();
103103
operator.depositIntoEigenlayer(strategies, tokenBalances);
104104

105-
// Roll passed the allocation configuration delay
105+
// Roll past the allocation configuration delay
106106
rollForward({blocks: ALLOCATION_CONFIGURATION_DELAY});
107-
}
108107

109-
assert_Snap_Added_Staker_DepositShares(operator, strategies, addedShares, "_newRandomOperator: failed to add delegatable shares");
108+
assert_Snap_Added_Staker_DepositShares(operator, strategies, addedShares, "_newRandomOperator: failed to add delegatable shares");
109+
}
110+
110111
assert_Snap_Added_OperatorShares(operator, strategies, addedShares, "_newRandomOperator: failed to award shares to operator");
111112
assertTrue(delegationManager.isOperator(address(operator)), "_newRandomOperator: operator should be registered");
112113

@@ -141,13 +142,13 @@ abstract contract IntegrationBase is IntegrationDeployer {
141142

142143
/// @dev Choose a random subset of validators (selects AT LEAST ONE)
143144
function _choose(uint40[] memory validators) internal returns (uint40[] memory) {
144-
uint rand = _randUint({ min: 1, max: validators.length ** 2 });
145+
uint _rand = _randUint({ min: 1, max: validators.length ** 2 });
145146

146147
uint40[] memory result = new uint40[](validators.length);
147148
uint newLen;
148149
for (uint i = 0; i < validators.length; i++) {
149150
// if bit set, add validator
150-
if (rand >> i & 1 == 1) {
151+
if (_rand >> i & 1 == 1) {
151152
result[newLen] = validators[i];
152153
newLen++;
153154
}
@@ -645,20 +646,15 @@ abstract contract IntegrationBase is IntegrationDeployer {
645646
function assert_Snap_Added_Staker_DepositShares(
646647
User staker,
647648
IStrategy strat,
648-
uint _addedShares,
649+
uint addedShares,
649650
string memory err
650651
) internal {
651-
IStrategy[] memory strategies = new IStrategy[](1);
652-
uint[] memory addedShares = new uint[](1);
653-
strategies[0] = strat;
654-
addedShares[0] = _addedShares;
655-
656-
assert_Snap_Added_Staker_DepositShares(staker, strategies, addedShares, err);
652+
assert_Snap_Added_Staker_DepositShares(staker, strat.toArray(), addedShares.toArrayU256(), err);
657653
}
658654

659655
/// @dev Check that the staker has `removedShares` fewer delegatable shares
660656
/// for each strategy since the last snapshot
661-
function assert_Snap_Removed_StakerDepositShares(
657+
function assert_Snap_Removed_Staker_DepositShares(
662658
User staker,
663659
IStrategy[] memory strategies,
664660
uint[] memory removedShares,
@@ -674,22 +670,52 @@ abstract contract IntegrationBase is IntegrationDeployer {
674670
}
675671
}
676672

677-
function assert_Snap_Removed_StakerDepositShares(
673+
function assert_Snap_Removed_Staker_DepositShares(
678674
User staker,
679675
IStrategy strat,
680-
uint _removedShares,
676+
uint removedShares,
677+
string memory err
678+
) internal {
679+
assert_Snap_Removed_Staker_DepositShares(staker, strat.toArray(), removedShares.toArrayU256(), err);
680+
}
681+
682+
/// @dev Check that the staker's delegatable shares in ALL strategies have not changed
683+
/// since the last snapshot
684+
function assert_Snap_Unchanged_Staker_DepositShares(
685+
User staker,
681686
string memory err
682687
) internal {
683-
IStrategy[] memory strategies = new IStrategy[](1);
684-
uint[] memory removedShares = new uint[](1);
685-
strategies[0] = strat;
686-
removedShares[0] = _removedShares;
688+
IStrategy[] memory strategies = allStrats;
689+
690+
uint[] memory curShares = _getStakerDepositShares(staker, strategies);
691+
// Use timewarp to get previous staker shares
692+
uint[] memory prevShares = _getPrevStakerDepositShares(staker, strategies);
687693

688-
assert_Snap_Removed_StakerDepositShares(staker, strategies, removedShares, err);
694+
// For each strategy, check (prev == cur)
695+
for (uint i = 0; i < strategies.length; i++) {
696+
assertEq(prevShares[i], curShares[i], err);
697+
}
689698
}
690699

691700
/// @dev Check that the staker's withdrawable shares have decreased by `removedShares`
692-
function assert_Snap_Removed_StakerWithdrawableShares(
701+
function assert_Snap_Added_Staker_WithdrawableShares(
702+
User staker,
703+
IStrategy[] memory strategies,
704+
uint[] memory addedShares,
705+
string memory err
706+
) internal {
707+
uint[] memory curShares = _getStakerWithdrawableShares(staker, strategies);
708+
// Use timewarp to get previous staker shares
709+
uint[] memory prevShares = _getPrevStakerWithdrawableShares(staker, strategies);
710+
711+
// For each strategy, check (prev - removed == cur)
712+
for (uint i = 0; i < strategies.length; i++) {
713+
assertEq(prevShares[i] + addedShares[i], curShares[i], err);
714+
}
715+
}
716+
717+
/// @dev Check that the staker's withdrawable shares have decreased by `removedShares`
718+
function assert_Snap_Removed_Staker_WithdrawableShares(
693719
User staker,
694720
IStrategy[] memory strategies,
695721
uint[] memory removedShares,
@@ -705,18 +731,30 @@ abstract contract IntegrationBase is IntegrationDeployer {
705731
}
706732
}
707733

708-
function assert_Snap_Removed_StakerWithdrawableShares(
734+
function assert_Snap_Removed_Staker_WithdrawableShares(
709735
User staker,
710736
IStrategy strat,
711-
uint _removedShares,
737+
uint removedShares,
738+
string memory err
739+
) internal {
740+
assert_Snap_Removed_Staker_WithdrawableShares(staker, strat.toArray(), removedShares.toArrayU256(), err);
741+
}
742+
743+
/// @dev Check that the staker's withdrawable shares have decreased by `removedShares`
744+
function assert_Snap_Unchanged_Staker_WithdrawableShares(
745+
User staker,
712746
string memory err
713747
) internal {
714-
IStrategy[] memory strategies = new IStrategy[](1);
715-
uint[] memory removedShares = new uint[](1);
716-
strategies[0] = strat;
717-
removedShares[0] = _removedShares;
748+
IStrategy[] memory strategies = allStrats;
749+
750+
uint[] memory curShares = _getStakerWithdrawableShares(staker, strategies);
751+
// Use timewarp to get previous staker shares
752+
uint[] memory prevShares = _getPrevStakerWithdrawableShares(staker, strategies);
718753

719-
assert_Snap_Removed_StakerWithdrawableShares(staker, strategies, removedShares, err);
754+
// For each strategy, check (prev - removed == cur)
755+
for (uint i = 0; i < strategies.length; i++) {
756+
assertEq(prevShares[i], curShares[i], err);
757+
}
720758
}
721759

722760
/// @dev Check that the staker's withdrawable shares have decreased by at least `removedShares`
@@ -743,30 +781,7 @@ abstract contract IntegrationBase is IntegrationDeployer {
743781
uint removedShares,
744782
string memory err
745783
) internal {
746-
IStrategy[] memory strategies = new IStrategy[](1);
747-
uint[] memory removedSharesArr = new uint[](1);
748-
strategies[0] = strat;
749-
removedSharesArr[0] = removedShares;
750-
751-
assert_Snap_Removed_Staker_WithdrawableShares_AtLeast(staker, strategies, removedSharesArr, err);
752-
}
753-
754-
/// @dev Check that the staker's delegatable shares in ALL strategies have not changed
755-
/// since the last snapshot
756-
function assert_Snap_Unchanged_StakerDepositShares(
757-
User staker,
758-
string memory err
759-
) internal {
760-
IStrategy[] memory strategies = allStrats;
761-
762-
uint[] memory curShares = _getStakerDepositShares(staker, strategies);
763-
// Use timewarp to get previous staker shares
764-
uint[] memory prevShares = _getPrevStakerDepositShares(staker, strategies);
765-
766-
// For each strategy, check (prev == cur)
767-
for (uint i = 0; i < strategies.length; i++) {
768-
assertEq(prevShares[i], curShares[i], err);
769-
}
784+
assert_Snap_Removed_Staker_WithdrawableShares_AtLeast(staker, strat.toArray(), removedShares.toArrayU256(), err);
770785
}
771786

772787
function assert_Snap_Delta_StakerShares(
@@ -1485,24 +1500,15 @@ abstract contract IntegrationBase is IntegrationDeployer {
14851500
// This method should only be used for tests that handle positive
14861501
// balances. Negative balances are an edge case that require
14871502
// the own tests and helper methods.
1488-
int shares;
1489-
if (forkType != LOCAL && !isUpgraded) {
1490-
shares = int(IEigenPodManager_DeprecatedM2(address(eigenPodManager)).podOwnerShares(address(staker)));
1491-
} else {
1492-
shares = int(eigenPodManager.podOwnerDepositShares(address(staker)));
1493-
}
1503+
int shares = eigenPodManager.podOwnerDepositShares(address(staker));
14941504

14951505
if (shares < 0) {
14961506
revert("_getStakerDepositShares: negative shares");
14971507
}
14981508

14991509
curShares[i] = uint(shares);
15001510
} else {
1501-
if (forkType != LOCAL && !isUpgraded) {
1502-
curShares[i] = IStrategyManager_DeprecatedM2(address(strategyManager)).stakerStrategyShares(address(staker), strat);
1503-
} else {
1504-
curShares[i] = strategyManager.stakerDepositShares(address(staker), strat);
1505-
}
1511+
curShares[i] = strategyManager.stakerDepositShares(address(staker), strat);
15061512
}
15071513
}
15081514

@@ -1656,13 +1662,8 @@ abstract contract IntegrationBase is IntegrationDeployer {
16561662
}
16571663

16581664
function _getCheckpointPodBalanceGwei(User staker) internal view returns (uint64) {
1659-
if (forkType != LOCAL && !isUpgraded) {
1660-
IEigenPod_DeprecatedM2 pod = IEigenPod_DeprecatedM2(address(staker.pod()));
1661-
return uint64(pod.currentCheckpoint().podBalanceGwei);
1662-
} else {
1663-
EigenPod pod = staker.pod();
1664-
return uint64(pod.currentCheckpoint().podBalanceGwei);
1665-
}
1665+
EigenPod pod = staker.pod();
1666+
return uint64(pod.currentCheckpoint().podBalanceGwei);
16661667
}
16671668

16681669
function _getPrevCheckpointPodBalanceGwei(User staker) internal timewarp() returns (uint64) {

0 commit comments

Comments
 (0)