From 6f5cb23c99964c50381f4746ea61672018b39716 Mon Sep 17 00:00:00 2001
From: lucasmt <lmartinellitabajara@gmail.com>
Date: Tue, 6 Aug 2024 15:14:28 -0500
Subject: [PATCH] Add stETH assumptions and invariants

---
 test/kontrol/StorageSetup.sol | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/test/kontrol/StorageSetup.sol b/test/kontrol/StorageSetup.sol
index c956d312..6ff8f2e2 100644
--- a/test/kontrol/StorageSetup.sol
+++ b/test/kontrol/StorageSetup.sol
@@ -32,6 +32,26 @@ contract StorageSetup is KontrolTest {
         _stEth.setShares(address(_escrow), shares);
     }
 
+    function stEthStorageInvariants(Mode mode, StETHModel _stEth, IEscrow _escrow) external {
+        uint256 totalPooledEther = _stEth.getTotalPooledEther();
+        uint256 totalShares = _stEth.getTotalShares();
+        uint256 escrowShares = _stEth.sharesOf(address(_escrow));
+
+        _establish(mode, 0 < _stEth.getTotalPooledEther());
+        _establish(mode, 0 < _stEth.getTotalShares());
+        _establish(mode, escrowShares < totalShares);
+    }
+
+    function stEthAssumeBounds(StETHModel _stEth, IEscrow _escrow) external {
+        uint256 totalPooledEther = _stEth.getTotalPooledEther();
+        uint256 totalShares = _stEth.getTotalShares();
+        uint256 escrowShares = _stEth.sharesOf(address(_escrow));
+
+        vm.assume(totalPooledEther < ethUpperBound);
+        vm.assume(totalShares < ethUpperBound);
+        vm.assume(escrowShares < ethUpperBound);
+    }
+
     function _wstEthStorageSetup(WstETHAdapted _wstEth, IStETH _stEth) internal {
         kevm.symbolicStorage(address(_wstEth));
     }