Skip to content

Commit

Permalink
test: refactor config rand
Browse files Browse the repository at this point in the history
  • Loading branch information
@wadealexc
wadealexc committed Nov 17, 2023
1 parent 02cc0f2 commit 6006085
Show file tree
Hide file tree
Showing 4 changed files with 324 additions and 238 deletions.
262 changes: 31 additions & 231 deletions src/test/integration/IntegrationBase.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,253 +9,32 @@ import "src/test/integration/IntegrationDeployer.t.sol";
import "src/test/integration/Global.t.sol";
import "src/test/integration/users/User.sol";

abstract contract Flags is Test {

uint constant FLAG = 1;

/// @dev Asset flags
/// These are used with _configRand to determine what assets are given
/// to a user when they are created.
uint constant NO_ASSETS = (FLAG << 0); // will have no assets
uint constant HOLDS_LST = (FLAG << 1); // will hold some random amount of LSTs
uint constant HOLDS_ETH = (FLAG << 2); // will hold some random amount of ETH
uint constant HOLDS_MIX = (FLAG << 3); // will hold a mix of LSTs and ETH

/// @dev Withdrawal flags
/// These are used with _configRand to determine how a user conducts a withdrawal
uint constant FULL_WITHDRAW_SINGLE = (FLAG << 0); // stakers will withdraw all assets using a single queued withdrawal
uint constant FULL_WITHDRAW_MULTI = (FLAG << 1); // stakers will withdraw all assets using multiple queued withdrawals
uint constant PART_WITHDRAW_SINGLE = (FLAG << 2); // stakers will withdraw some, but not all assets

/// Note: Thought about the following flags (but did not implement) -
///
/// WithdrawerType (SELF_WITHDRAWER, OTHER_WITHDRAWER)
/// - especially with EPM share handling, this felt like it deserved its own test rather than a fuzzy state
/// CompletionType (AS_TOKENS, AS_SHARES)
/// - same reason as above
///
/// DepositMethod (DEPOSIT_STD, DEPOSIT_WITH_SIG)
/// - could still do this!
/// WithdrawalMethod (QUEUE_WITHDRAWAL, UNDELEGATE, REDELEGATE)
/// - could still do this!
/// - This would trigger staker.queueWithdrawals to use either `queueWithdrawals` or `undelegate` under the hood
/// - "redelegate" would be like the above, but adding a new `delegateTo` step after undelegating

bytes32 random;

bytes stakerAssets;
bytes operatorAssets;
bytes withdrawTypes;

function _configRand(
uint16 _randomSeed,
uint _stakerAssets,
uint _operatorAssets,
uint _withdrawType
) internal {
// Using uint16 for the seed type so that if a test fails, it's easier
// to manually use the seed to replay the same test.
// TODO - should we expand the type?
emit log_named_uint("_configRand: set random seed to ", _randomSeed);
random = keccak256(abi.encodePacked(_randomSeed));

emit log_named_uint("staker assets: ", _stakerAssets);
emit log_named_uint("operator assets: ", _operatorAssets);

emit log_named_uint("HOLDS_LST: ", HOLDS_LST);

// Convert flag bitmaps to bytes of set bits for easy use with _randUint
stakerAssets = _bitmapToBytes(_stakerAssets);
operatorAssets = _bitmapToBytes(_operatorAssets);
withdrawTypes = _bitmapToBytes(_withdrawType);

emit log("staker assets");
for (uint i = 0; i < stakerAssets.length; i++) {
emit log_named_uint("- ", uint(uint8(stakerAssets[i])));
}

emit log("operator assets");
for (uint i = 0; i < operatorAssets.length; i++) {
emit log_named_uint("- ", uint(uint8(operatorAssets[i])));
}

assertTrue(stakerAssets.length != 0, "_configRand: no staker assets selected");
assertTrue(operatorAssets.length != 0, "_configRand: no operator assets selected");
assertTrue(withdrawTypes.length != 0, "_configRand: no withdrawal type selected");
}
abstract contract IntegrationBase is IntegrationDeployer {

/**
* @dev Converts a bitmap into an array of bytes
* @dev Each byte in the input is processed as indicating a single bit to flip in the bitmap
* Gen/Init methods:
*/
function _bitmapToBytes(uint bitmap) internal returns (bytes memory bytesArray) {
for (uint i = 0; i < 256; ++i) {
// Mask for i-th bit
uint mask = uint(1 << i);

// emit log_named_uint("mask: ", mask);

// If the i-th bit is flipped, add a byte to the return array
if (bitmap & mask != 0) {
emit log_named_uint("i: ", i);
bytesArray = bytes.concat(bytesArray, bytes1(uint8(1 << i)));
}
}
return bytesArray;
}

/// @dev Uses `random` to return a random uint, with a range given by `min` and `max` (inclusive)
/// @return `min` <= result <= `max`
function _randUint(uint min, uint max) internal returns (uint) {
uint range = max - min + 1;

// calculate the number of bits needed for the range
uint bitsNeeded = 0;
uint tempRange = range;
while (tempRange > 0) {
bitsNeeded++;
tempRange >>= 1;
}

// create a mask for the required number of bits
// and extract the value from the hash
uint mask = (1 << bitsNeeded) - 1;
uint value = uint(random) & mask;

// in case value is out of range, wrap around or retry
while (value >= range) {
value = (value - range) & mask;
}

// Hash `random` with itself so the next value we generate is different
random = keccak256(abi.encodePacked(random));
return min + value;
}
}

abstract contract IntegrationBase is IntegrationDeployer, Flags {

Global global;

function setUp() public virtual override {
super.setUp();

global = new Global();
}

/**
* Gen/Init methods:
* @dev Create a new user according to configured random variants.
* This user is ready to deposit into some strategies and has some underlying token balances
*/

/// @dev Uses configured randomness to create a Staker-type user
/// @return The new User, the strategies the user has token balances for, and the
/// token balances themselves.
function _newStaker() internal returns (User, IStrategy[] memory, uint[] memory) {
// Select assets for the staker
uint assetIdx = _randUint({ min: 0, max: stakerAssets.length - 1 });
emit log_named_uint("_newStaker: rand idx: ", assetIdx);

uint flag = uint(uint8(stakerAssets[assetIdx]));

emit log_named_uint("_newStaker: asset flag is ", flag);

User staker = new User(delegationManager, strategyManager, eigenPodManager, global);
IStrategy[] memory strategies;
uint[] memory tokenBalances;

if (flag == NO_ASSETS) {
emit log("_newStaker: creating NO_ASSETS user");
} else if (flag == HOLDS_LST) {
emit log("_newStaker: creating HOLDS_LST user");

strategies = new IStrategy[](1);
tokenBalances = new uint[](1);

// TODO rand more than 1 strategy
uint idx = _randUint({ min: 0, max: _strategies.length - 1 });

IStrategy strategy = _strategies[idx];
IERC20 underlyingToken = strategy.underlyingToken();
(User staker, IStrategy[] memory strategies, uint[] memory tokenBalances) = _randUser();

// TODO better rand token balance
uint balance = _randUint({ min: 1e6, max: 5e6 });

// award tokens
StdCheats.deal(address(underlyingToken), address(staker), balance);

strategies[0] = strategy;
tokenBalances[0] = balance;

emit log_named_string("_newStaker: user receiving token: ", IERC20Metadata(address(underlyingToken)).name());
emit log_named_uint("_newStaker: awarded tokens to user: ", balance);
} else if (flag == HOLDS_ETH) {
emit log("_newStaker: creating HOLDS_ETH user");
// Need to support native eth initialization
revert("unimplemented");
} else if (flag == HOLDS_MIX) {
emit log("_newStaker: creating HOLDS_MIX user");
// Need to support native eth initialization
revert("unimplemented");
} else {
revert("_newStaker: invalid flag");
}
assert_HasUnderlyingTokenBalances(staker, strategies, tokenBalances, "_newStaker: failed to award token balances");

return (staker, strategies, tokenBalances);
}

function _newOperator() internal returns (User, IStrategy[] memory, uint[] memory) {
// Select assets for the operator
uint assetIdx = _randUint({ min: 0, max: operatorAssets.length - 1 });
emit log_named_uint("_newOperator: rand idx: ", assetIdx);

uint flag = uint(uint8(operatorAssets[assetIdx]));

emit log_named_uint("_newOperator: asset flag is ", flag);

User operator = new User(delegationManager, strategyManager, eigenPodManager, global);
IStrategy[] memory strategies;
uint[] memory tokenBalances;

if (flag == NO_ASSETS) {
emit log("_newOperator: creating NO_ASSETS user");
} else if (flag == HOLDS_LST) {
emit log("_newOperator: creating HOLDS_LST user");

strategies = new IStrategy[](1);
tokenBalances = new uint[](1);

// TODO rand more than 1 strategy
uint idx = _randUint({ min: 0, max: _strategies.length - 1 });

IStrategy strategy = _strategies[idx];
IERC20 underlyingToken = strategy.underlyingToken();
(User operator, IStrategy[] memory strategies, uint[] memory tokenBalances) = _randUser();

// TODO better rand token balance
uint balance = _randUint({ min: 1e6, max: 5e6 });

// award tokens
StdCheats.deal(address(underlyingToken), address(operator), balance);

strategies[0] = strategy;
tokenBalances[0] = balance;

emit log_named_string("_newOperator: user receiving token: ", IERC20Metadata(address(underlyingToken)).name());
emit log_named_uint("_newOperator: awarded tokens to user: ", balance);
} else if (flag == HOLDS_ETH) {
emit log("_newOperator: creating HOLDS_ETH user");
// Need to support native eth initialization
revert("unimplemented");
} else if (flag == HOLDS_MIX) {
emit log("_newOperator: creating HOLDS_MIX user");
// Need to support native eth initialization
revert("unimplemented");
} else {
revert("_newOperator: invalid flag");
}

operator.registerAsOperator();
// TODO - deposit into strats

assert_HasUnderlyingTokenBalances(operator, strategies, tokenBalances, "_newOperator: failed to award token balances");
assertTrue(delegationManager.isOperator(address(operator)), "_newOperator: operator should be registered");
// assert_HasNoDelegatableShares(operator, "_newOperator: new operator should not have delegatable shares");

return (operator, strategies, tokenBalances);
}
Expand All @@ -272,6 +51,27 @@ abstract contract IntegrationBase is IntegrationDeployer, Flags {
assertEq(strategies.length, shares.length, "assert_HasNoDelegatableShares: return length mismatch");
}

function assert_HasUnderlyingTokenBalances(
User user,
IStrategy[] memory strategies,
uint[] memory expectedBalances,
string memory err
) internal {
for (uint i = 0; i < strategies.length; i++) {
IStrategy strat = strategies[i];

uint expectedBalance = expectedBalances[i];
uint tokenBalance;
if (strat == BEACONCHAIN_ETH_STRAT) {
tokenBalance = address(user).balance;
} else {
tokenBalance = strat.underlyingToken().balanceOf(address(user));
}

assertEq(expectedBalance, tokenBalance, err);
}
}

function assert_HasNoUnderlyingTokenBalance(User user, IStrategy[] memory strategies, string memory err) internal {
for (uint i = 0; i < strategies.length; i++) {
IStrategy strat = strategies[i];
Expand Down
Loading

0 comments on commit 6006085

Please sign in to comment.