diff --git a/src/test/integration/IntegrationBase.t.sol b/src/test/integration/IntegrationBase.t.sol index cc569cebfa..c139396d8d 100644 --- a/src/test/integration/IntegrationBase.t.sol +++ b/src/test/integration/IntegrationBase.t.sol @@ -314,6 +314,48 @@ abstract contract IntegrationBase is IntegrationDeployer, TypeImporter { assertEq(withdrawalRoot, delegationManager.calculateWithdrawalRoot(withdrawal), err); } + function assert_StakerStrategyListEmpty( + User staker, + string memory err + ) internal view { + IStrategy[] memory strategies = _getStakerStrategyList(staker); + assertEq(strategies.length, 0, err); + } + + function assert_StrategyNotInStakerStrategyList( + User staker, + IStrategy strategy, + string memory err + ) internal view { + // BEACONCHAIN_ETH_STRAT is not in the staker's strategy list + if (strategy == BEACONCHAIN_ETH_STRAT) return; + + IStrategy[] memory strategies = _getStakerStrategyList(staker); + assertFalse(strategies.contains(strategy), err); + } + + function assert_StrategiesInStakerStrategyList( + User staker, + IStrategy[] memory strategies, + string memory err + ) internal view { + for(uint i = 0; i < strategies.length; i++) { + assert_StrategyInStakerStrategyList(staker, strategies[i], err); + } + } + + function assert_StrategyInStakerStrategyList( + User staker, + IStrategy strategy, + string memory err + ) internal view { + // BEACONCHAIN_ETH_STRAT is not in the staker's strategy list + if (strategy == BEACONCHAIN_ETH_STRAT) return; + + IStrategy[] memory strategies = _getStakerStrategyList(staker); + assertTrue(strategies.contains(strategy), err); + } + function assert_PodBalance_Eq( User staker, uint expectedBalance, @@ -3058,6 +3100,10 @@ abstract contract IntegrationBase is IntegrationDeployer, TypeImporter { return curShares; } + function _getStakerStrategyList(User staker) internal view returns (IStrategy[] memory) { + return strategyManager.getStakerStrategyList(address(staker)); + } + function _getPrevStakerWithdrawableShares(User staker, IStrategy[] memory strategies) internal timewarp() returns (uint[] memory) { return _getStakerWithdrawableShares(staker, strategies); } diff --git a/src/test/integration/IntegrationChecks.t.sol b/src/test/integration/IntegrationChecks.t.sol index c88eb071b9..b93c92bf16 100644 --- a/src/test/integration/IntegrationChecks.t.sol +++ b/src/test/integration/IntegrationChecks.t.sol @@ -201,6 +201,7 @@ contract IntegrationCheckUtils is IntegrationBase { // 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_Added_Staker_DepositShares(staker, strategies, shares, "staker should expect shares in each strategy after depositing"); + assert_StrategiesInStakerStrategyList(staker, strategies, "staker strategy list should contain all strategies"); if (delegationManager.isDelegated(address(staker))) { User operator = User(payable(delegationManager.delegatedTo(address(staker)))); @@ -219,6 +220,7 @@ contract IntegrationCheckUtils is IntegrationBase { // and that the staker now has the expected amount of delegated shares in each strategy assert_HasUnderlyingTokenBalances(staker, strategies, tokenBalances, "staker should have transferred some underlying tokens"); assert_Snap_Added_Staker_DepositShares(staker, strategies, shares, "staker should expected shares in each strategy after depositing"); + assert_StrategiesInStakerStrategyList(staker, strategies, "staker strategy list should contain all strategies"); if (delegationManager.isDelegated(address(staker))) { User operator = User(payable(delegationManager.delegatedTo(address(staker)))); @@ -295,15 +297,17 @@ contract IntegrationCheckUtils is IntegrationBase { check_Decreased_SlashableStake(operator, withdrawableShares, strategies); // Check that the dsf is either reset to wad or unchanged for (uint i = 0; i < strategies.length; i++) { - // For a full withdrawal, the dsf should be reset to wad + // For a full withdrawal, the dsf should be reset to wad & the staker strategy list should not contain the strategy if (_getStakerDepositShares(staker, strategies[i].toArray())[0] == 0) { assert_DSF_Reset(staker, strategies[i].toArray(), "check_QueuedWithdrawal_State: dsf should be reset to wad"); + assert_StrategyNotInStakerStrategyList(staker, strategies[i], "check_QueuedWithdrawal_State: staker strategy list should not contain strategy"); } - // For a partial withdrawal, the dsf should not be changed + // For a partial withdrawal, the dsf should not be changed & the strategy should still be in the staker strategy list else { assert_Snap_Unchanged_DSF(staker, strategies[i].toArray(), "check_QueuedWithdrawal_State: dsf should not be changed"); + assert_StrategyInStakerStrategyList(staker, strategies[i], "check_QueuedWithdrawal_State: staker strategy list should contain strategy"); } } } @@ -347,13 +351,14 @@ contract IntegrationCheckUtils is IntegrationBase { "check_Undelegate_State: stakers withdrawal should now be pending"); assert_DSF_Reset(staker, strategies, "check_Undelegate_State: staker dsfs should be reset to wad"); + assert_StakerStrategyListEmpty(staker, "check_Undelegate_State: staker strategy list should be empty"); assert_Snap_Added_QueuedWithdrawals(staker, withdrawals, "check_Undelegate_State: staker should have increased nonce by withdrawals.length"); assert_Snap_Removed_OperatorShares(operator, strategies, stakerDelegatedShares, "check_Undelegate_State: failed to remove operator shares"); assert_RemovedAll_Staker_DepositShares(staker, strategies, "check_Undelegate_State: failed to remove staker shares"); assert_RemovedAll_Staker_WithdrawableShares(staker, strategies, - "check_QueuedWithdrawal_State: failed to remove staker withdrawable shares"); + "check_Undelegate_State: failed to remove staker withdrawable shares"); } function check_Redelegate_State( @@ -381,11 +386,12 @@ contract IntegrationCheckUtils is IntegrationBase { "check_Redelegate_State: stakers withdrawal should now be pending"); assert_Snap_Added_QueuedWithdrawals(staker, withdrawals, "check_Redelegate_State: staker should have increased nonce by withdrawals.length"); + assert_StakerStrategyListEmpty(staker, "check_Redelegate_State: staker strategy list should be empty"); assert_Snap_Removed_OperatorShares(oldOperator, strategies, stakerDelegatedShares, "check_Redelegate_State: failed to remove operator shares"); assert_RemovedAll_Staker_DepositShares(staker, strategies, "check_Undelegate_State: failed to remove staker shares"); assert_RemovedAll_Staker_WithdrawableShares(staker, strategies, - "check_QueuedWithdrawal_State: failed to remove staker withdrawable shares"); + "check_Redelegate_State: failed to remove staker withdrawable shares"); assert_Snap_Unchanged_OperatorShares(newOperator, "check_Redelegate_State: new operator shares should not have changed"); }