add invariant measurement

yieldnest · Sep 1, 2024 · a4cf954 · a4cf954
1 parent 014cecc
commit a4cf954
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 0 deletions.
diff --git a/test/integration/M3/Base.t.sol b/test/integration/M3/Base.t.sol
@@ -27,6 +27,8 @@ import {RewardsDistributor} from "../../../src/RewardsDistributor.sol";
 import {StakingNode} from "../../../src/StakingNode.sol";
 import {WithdrawalQueueManager} from "../../../src/WithdrawalQueueManager.sol";
 import {ynETHRedemptionAssetsVault} from "../../../src/ynETHRedemptionAssetsVault.sol";
+import {IStakingNode} from "../../../src/interfaces/IStakingNodesManager.sol";
+
 
 import "forge-std/console.sol";
 import "forge-std/Test.sol";
@@ -257,4 +259,28 @@ contract Base is Test, Utils {
         vm.prank(actors.ops.VALIDATOR_MANAGER);
         stakingNodesManager.registerValidators(validatorData);
     }
+
+    function runSystemStateInvariants(
+        uint256 previousTotalAssets,
+        uint256 previousTotalSupply,
+        uint256[] memory previousStakingNodeBalances
+    ) public {  
+        assertEq(yneth.totalAssets(), previousTotalAssets, "Total assets integrity check failed");
+        assertEq(yneth.totalSupply(), previousTotalSupply, "Share mint integrity check failed");
+        for (uint i = 0; i < previousStakingNodeBalances.length; i++) {
+            IStakingNode stakingNodeInstance = stakingNodesManager.nodes(i);
+            uint256 currentStakingNodeBalance = stakingNodeInstance.getETHBalance();
+            assertEq(currentStakingNodeBalance, previousStakingNodeBalances[i], "Staking node balance integrity check failed for node ID: ");
+        }
+	}
+
+    function getAllStakingNodeBalances() public view returns (uint256[] memory) {
+        uint256[] memory balances = new uint256[](stakingNodesManager.nodesLength());
+        for (uint256 i = 0; i < stakingNodesManager.nodesLength(); i++) {
+            IStakingNode stakingNode = stakingNodesManager.nodes(i);
+            balances[i] = stakingNode.getETHBalance();
+        }
+        return balances;
+    }
+
 }
diff --git a/test/integration/M3/WithdrawalsWithRewards-Scenario.t.sol b/test/integration/M3/WithdrawalsWithRewards-Scenario.t.sol
@@ -29,6 +29,15 @@ contract M3WithdrawalsTest is Base {
 
     uint256 public amount;
 
+    struct TestState {
+        uint256 totalAssetsBefore;
+        uint256 totalSupplyBefore;
+        uint256[] stakingNodeBalancesBefore;
+        uint256 previousYnETHRedemptionAssetsVaultBalance;
+        uint256 previousYnETHBalance;
+    }
+
+
     function setUp() public override {
         super.setUp();
     }
@@ -47,6 +56,15 @@ contract M3WithdrawalsTest is Base {
             uint256 userYnETHBalance = yneth.balanceOf(user);
             console.log("User ynETH balance after deposit:", userYnETHBalance);
         }
+
+        TestState memory state = TestState({
+            totalAssetsBefore: yneth.totalAssets(),
+            totalSupplyBefore: yneth.totalSupply(),
+            stakingNodeBalancesBefore: getAllStakingNodeBalances(),
+            previousYnETHRedemptionAssetsVaultBalance: ynETHRedemptionAssetsVaultInstance.availableRedemptionAssets(),
+            previousYnETHBalance: address(yneth).balance
+        });
+
         // create staking node
         {
             vm.prank(actors.ops.STAKING_NODE_CREATOR);
@@ -73,6 +91,8 @@ contract M3WithdrawalsTest is Base {
             registerValidators(nodeIds);
         }
 
+        // runSystemStateInvariants(state.totalAssetsBefore, state.totalSupplyBefore, state.stakingNodeBalancesBefore);
+
         // verify withdrawal credentials
         {