Skip to content

Commit

Permalink
add state assertions
Browse files Browse the repository at this point in the history
  • Loading branch information
@danoctavian
danoctavian committed Oct 1, 2024
1 parent 0d5a1ba commit 5844c3d
Showing 1 changed file with 116 additions and 79 deletions.
195 changes: 116 additions & 79 deletions test/integration/StakingNode.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -318,6 +318,25 @@ contract StakingNodeVerifyWithdrawalCredentials is StakingNodeTestBase {

yneth.depositETH{value: 1000 ether}(user);
}
struct StateSnapshot {
uint256 totalAssets;
uint256 totalSupply;
uint256 stakingNodeBalance;
uint256 queuedShares;
uint256 withdrawnETH;
uint256 unverifiedStakedETH;
}

function takeSnapshot(uint256 nodeId) internal view returns (StateSnapshot memory) {
return StateSnapshot({
totalAssets: yneth.totalAssets(),
totalSupply: yneth.totalSupply(),
stakingNodeBalance: stakingNodesManager.nodes(nodeId).getETHBalance(),
queuedShares: stakingNodesManager.nodes(nodeId).getQueuedSharesAmount(),
withdrawnETH: stakingNodesManager.nodes(nodeId).getWithdrawnETH(),
unverifiedStakedETH: stakingNodesManager.nodes(nodeId).unverifiedStakedETH()
});
}

function testVerifyWithdrawalCredentialsForOneValidator() public {

Expand All @@ -328,37 +347,27 @@ contract StakingNodeVerifyWithdrawalCredentials is StakingNodeTestBase {
registerValidators(repeat(nodeId, 1));


// Capture state before verification
uint256 totalAssetsBefore = yneth.totalAssets();
uint256 totalSupplyBefore = yneth.totalSupply();
uint256 stakingNodeBalanceBefore = stakingNodesManager.nodes(nodeId).getETHBalance();
uint256 queuedSharesBefore = stakingNodesManager.nodes(nodeId).queuedShares();
uint256 withdrawnETHBefore = stakingNodesManager.nodes(nodeId).withdrawnETH();
uint256 unverifiedStakedETHBefore = stakingNodesManager.nodes(nodeId).unverifiedStakedETH();

_verifyWithdrawalCredentials(nodeId, validatorIndices[0]);

// Capture state after verification
uint256 totalAssetsAfter = yneth.totalAssets();
uint256 totalSupplyAfter = yneth.totalSupply();
uint256 stakingNodeBalanceAfter = stakingNodesManager.nodes(nodeId).getETHBalance();
uint256 queuedSharesAfter = stakingNodesManager.nodes(nodeId).queuedShares();
uint256 withdrawnETHAfter = stakingNodesManager.nodes(nodeId).withdrawnETH();
uint256 unverifiedStakedETHAfter = stakingNodesManager.nodes(nodeId).unverifiedStakedETH();

// Assert that ynETH totalAssets, totalSupply, and staking Node balance, queuedShares and withdrawnETH stay the same
assertEq(totalAssetsAfter, totalAssetsBefore, "Total assets should not change");
assertEq(totalSupplyAfter, totalSupplyBefore, "Total supply should not change");
assertEq(stakingNodeBalanceAfter, stakingNodeBalanceBefore, "Staking node balance should not change");
assertEq(queuedSharesAfter, queuedSharesBefore, "Queued shares should not change");
assertEq(withdrawnETHAfter, withdrawnETHBefore, "Withdrawn ETH should not change");

// Assert that unverifiedStakedETH decreases
assertLt(unverifiedStakedETHAfter, unverifiedStakedETHBefore, "Unverified staked ETH should decrease");

// Additional checks
assertEq(unverifiedStakedETHAfter, 0, "Unverified staked ETH should be 0 after verification");
assertEq(uint256(eigenPodManager.podOwnerShares(address(stakingNodesManager.nodes(nodeId)))), AMOUNT, "Pod owner shares should equal AMOUNT");
// Capture state before verification
StateSnapshot memory before = takeSnapshot(nodeId);

_verifyWithdrawalCredentials(nodeId, validatorIndices[0]);

// Capture state after verification
StateSnapshot memory afterVerification = takeSnapshot(nodeId);

// Assert that ynETH totalAssets, totalSupply, and staking Node balance, queuedShares and withdrawnETH stay the same
assertEq(afterVerification.totalAssets, before.totalAssets, "Total assets should not change");
assertEq(afterVerification.totalSupply, before.totalSupply, "Total supply should not change");
assertEq(afterVerification.stakingNodeBalance, before.stakingNodeBalance, "Staking node balance should not change");
assertEq(afterVerification.queuedShares, before.queuedShares, "Queued shares should not change");
assertEq(afterVerification.withdrawnETH, before.withdrawnETH, "Withdrawn ETH should not change");

// Assert that unverifiedStakedETH decreases
assertLt(afterVerification.unverifiedStakedETH, before.unverifiedStakedETH, "Unverified staked ETH should decrease");

// Additional checks
assertEq(afterVerification.unverifiedStakedETH, 0, "Unverified staked ETH should be 0 after verification");
assertEq(uint256(eigenPodManager.podOwnerShares(address(stakingNodesManager.nodes(nodeId)))), AMOUNT, "Pod owner shares should equal AMOUNT");
}

function testVerifyWithdrawalCredentialsTwice() public {
Expand Down Expand Up @@ -391,56 +400,84 @@ contract StakingNodeVerifyWithdrawalCredentials is StakingNodeTestBase {
vm.stopPrank();
}

function testVerifyCheckpointsForOneValidator() public {
// function testVerifyCheckpointsForOneValidator() public {

uint256 nodeId = createStakingNodes(1)[0];
// Call createValidators with the nodeIds array and validatorCount
validatorIndices = createValidators(repeat(nodeId, 1), 1);
beaconChain.advanceEpoch_NoRewards();
registerValidators(repeat(nodeId, 1));
// uint256 nodeId = createStakingNodes(1)[0];
// // Call createValidators with the nodeIds array and validatorCount
// validatorIndices = createValidators(repeat(nodeId, 1), 1);
// beaconChain.advanceEpoch_NoRewards();
// registerValidators(repeat(nodeId, 1));

uint40 validatorIndex = validatorIndices[0];

{
_verifyWithdrawalCredentials(nodeId, validatorIndex);

// check that unverifiedStakedETH is 0 and podOwnerShares is 32 ETH (AMOUNT)
assertEq(stakingNodesManager.nodes(nodeId).unverifiedStakedETH(), 0, "_testVerifyWithdrawalCredentials: E0");
assertEq(uint256(eigenPodManager.podOwnerShares(address(stakingNodesManager.nodes(nodeId)))), AMOUNT, "_testVerifyWithdrawalCredentials: E1");
}

beaconChain.advanceEpoch();

// start checkpoint
{
vm.startPrank(actors.ops.STAKING_NODES_OPERATOR);
stakingNodesManager.nodes(nodeId).startCheckpoint(true);
vm.stopPrank();

// make sure startCheckpoint cant be called again, which means that the checkpoint has started
IStakingNode _node = stakingNodesManager.nodes(nodeId);
vm.expectRevert("EigenPod._startCheckpoint: must finish previous checkpoint before starting another");
vm.prank(actors.ops.STAKING_NODES_OPERATOR);
_node.startCheckpoint(true);
}

// verify checkpoints
{
uint40[] memory _validators = new uint40[](1);
_validators[0] = validatorIndex;
IStakingNode _node = stakingNodesManager.nodes(nodeId);
CheckpointProofs memory _cpProofs = beaconChain.getCheckpointProofs(_validators, _node.eigenPod().currentCheckpointTimestamp());
IEigenPodSimplified(address(_node.eigenPod())).verifyCheckpointProofs({
balanceContainerProof: _cpProofs.balanceContainerProof,
proofs: _cpProofs.balanceProofs
});

// check that proofsRemaining is 0 and podOwnerShares is still 32 ETH (AMOUNT)
IEigenPod.Checkpoint memory _checkpoint = stakingNodesManager.nodes(nodeId).eigenPod().currentCheckpoint();
assertEq(_checkpoint.proofsRemaining, 0, "_testVerifyCheckpointsBeforeWithdrawalRequest: E0");
assertApproxEqAbs(uint256(eigenPodManager.podOwnerShares(address(stakingNodesManager.nodes(nodeId)))), AMOUNT, 1000000000, "_testVerifyCheckpointsBeforeWithdrawalRequest: E1");
}
}
// uint40 validatorIndex = validatorIndices[0];

// {
// _verifyWithdrawalCredentials(nodeId, validatorIndex);

// // check that unverifiedStakedETH is 0 and podOwnerShares is 32 ETH (AMOUNT)
// assertEq(stakingNodesManager.nodes(nodeId).unverifiedStakedETH(), 0, "_testVerifyWithdrawalCredentials: E0");
// assertEq(uint256(eigenPodManager.podOwnerShares(address(stakingNodesManager.nodes(nodeId)))), AMOUNT, "_testVerifyWithdrawalCredentials: E1");
// }

// beaconChain.advanceEpoch();
// // Capture initial state
// uint256 initialTotalAssets = yneth.totalAssets();
// uint256 initialTotalSupply = yneth.totalSupply();
// uint256 initialNodeBalance = stakingNodesManager.nodes(nodeId).getETHBalance();
// uint256 initialQueuedShares = stakingNodesManager.nodes(nodeId).getQueuedSharesAmount();
// uint256 initialWithdrawnETH = stakingNodesManager.nodes(nodeId).getWithdrawnETH();
// uint256 initialUnverifiedStakedETH = stakingNodesManager.nodes(nodeId).unverifiedStakedETH();
// uint256 initialPodOwnerShares = uint256(eigenPodManager.podOwnerShares(address(stakingNodesManager.nodes(nodeId))));

// // start checkpoint
// {
// vm.startPrank(actors.ops.STAKING_NODES_OPERATOR);
// stakingNodesManager.nodes(nodeId).startCheckpoint(true);
// vm.stopPrank();

// // make sure startCheckpoint cant be called again, which means that the checkpoint has started
// IStakingNode _node = stakingNodesManager.nodes(nodeId);
// vm.expectRevert("EigenPod._startCheckpoint: must finish previous checkpoint before starting another");
// vm.prank(actors.ops.STAKING_NODES_OPERATOR);
// _node.startCheckpoint(true);
// }

// // Assert that state remains unchanged after starting checkpoint
// assertEq(yneth.totalAssets(), initialTotalAssets, "Total assets changed after starting checkpoint");
// assertEq(yneth.totalSupply(), initialTotalSupply, "Total supply changed after starting checkpoint");
// assertEq(stakingNodesManager.nodes(nodeId).getETHBalance(), initialNodeBalance, "Node balance changed after starting checkpoint");
// assertEq(stakingNodesManager.nodes(nodeId).getQueuedSharesAmount(), initialQueuedShares, "Queued shares changed after starting checkpoint");
// assertEq(stakingNodesManager.nodes(nodeId).getWithdrawnETH(), initialWithdrawnETH, "Withdrawn ETH changed after starting checkpoint");
// assertEq(stakingNodesManager.nodes(nodeId).unverifiedStakedETH(), initialUnverifiedStakedETH, "Unverified staked ETH changed after starting checkpoint");

// // verify checkpoints
// {
// uint40[] memory _validators = new uint40[](1);
// _validators[0] = validatorIndex;
// IStakingNode _node = stakingNodesManager.nodes(nodeId);
// CheckpointProofs memory _cpProofs = beaconChain.getCheckpointProofs(_validators, _node.eigenPod().currentCheckpointTimestamp());
// IEigenPodSimplified(address(_node.eigenPod())).verifyCheckpointProofs({
// balanceContainerProof: _cpProofs.balanceContainerProof,
// proofs: _cpProofs.balanceProofs
// });

// // check that proofsRemaining is 0
// IEigenPod.Checkpoint memory _checkpoint = stakingNodesManager.nodes(nodeId).eigenPod().currentCheckpoint();
// assertEq(_checkpoint.proofsRemaining, 0, "_testVerifyCheckpointsBeforeWithdrawalRequest: E0");

// // Assert that node balance and shares increased by the amount of rewards
// uint256 newPodOwnerShares = uint256(eigenPodManager.podOwnerShares(address(stakingNodesManager.nodes(nodeId))));
// uint256 rewardsAmount = newPodOwnerShares - initialPodOwnerShares;
// assertGt(rewardsAmount, 0, "No rewards were added");
// assertEq(stakingNodesManager.nodes(nodeId).getETHBalance(), initialNodeBalance + rewardsAmount, "Node balance did not increase by rewards amount");

// // Assert that other state variables remain unchanged
// assertEq(yneth.totalAssets(), initialTotalAssets, "Total assets changed after verification");
// assertEq(yneth.totalSupply(), initialTotalSupply, "Total supply changed after verification");
// assertEq(stakingNodesManager.nodes(nodeId).getQueuedSharesAmount(), initialQueuedShares, "Queued shares changed after verification");
// assertEq(stakingNodesManager.nodes(nodeId).getWithdrawnETH(), initialWithdrawnETH, "Withdrawn ETH changed after verification");
// assertEq(stakingNodesManager.nodes(nodeId).unverifiedStakedETH(), initialUnverifiedStakedETH, "Unverified staked ETH changed after verification");
// }
// }


function testVerifyCheckpointsForManyValidators() public {
Expand Down

0 comments on commit 5844c3d

Please sign in to comment.