From 03f7acf2ccbc3c6b505a67d8c0b2ab4fdd1c1e36 Mon Sep 17 00:00:00 2001
From: lucasmt <lmartinellitabajara@gmail.com>
Date: Mon, 15 Jul 2024 10:28:08 -0500
Subject: [PATCH] Add batchesQueue to storage setup

---
 test/kontrol/DualGovernanceSetUp.sol | 22 ++++++++++++----------
 test/kontrol/StorageSetup.sol        | 17 +++++++++++++++--
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/test/kontrol/DualGovernanceSetUp.sol b/test/kontrol/DualGovernanceSetUp.sol
index bbbdf1ee..c28f9720 100644
--- a/test/kontrol/DualGovernanceSetUp.sol
+++ b/test/kontrol/DualGovernanceSetUp.sol
@@ -64,19 +64,21 @@ contract DualGovernanceSetUp is StorageSetup {
         // ?WORD9: claimedETH
         // ?WORD10: unfinalizedShares
         // ?WORD11: finalizedETH
-        // ?WORD12: rageQuitExtensionDelay
-        // ?WORD13: rageQuitWithdrawalsTimelock
-        // ?WORD14: rageQuitTimelockStartedAt
+        // ?WORD12: batchesQueue
+        // ?WORD13: rageQuitExtensionDelay
+        // ?WORD14: rageQuitWithdrawalsTimelock
+        // ?WORD15: rageQuitTimelockStartedAt
         _signallingEscrowStorageSetup(signallingEscrow, dualGovernance);
 
         // ?STORAGE2
-        // ?WORD15: lockedShares
-        // ?WORD16: claimedETH
-        // ?WORD17: unfinalizedShares
-        // ?WORD18: finalizedETH
-        // ?WORD19: rageQuitExtensionDelay
-        // ?WORD20: rageQuitWithdrawalsTimelock
-        // ?WORD21: rageQuitTimelockStartedAt
+        // ?WORD16: lockedShares
+        // ?WORD17: claimedETH
+        // ?WORD18: unfinalizedShares
+        // ?WORD19: finalizedETH
+        // ?WORD20: batchesQueue
+        // ?WORD21: rageQuitExtensionDelay
+        // ?WORD22: rageQuitWithdrawalsTimelock
+        // ?WORD23: rageQuitTimelockStartedAt
         _rageQuitEscrowStorageSetup(rageQuitEscrow, dualGovernance);
 
         // ?STORAGE3
diff --git a/test/kontrol/StorageSetup.sol b/test/kontrol/StorageSetup.sol
index af408cd7..168ecf1c 100644
--- a/test/kontrol/StorageSetup.sol
+++ b/test/kontrol/StorageSetup.sol
@@ -5,6 +5,7 @@ import "contracts/EmergencyProtectedTimelock.sol";
 import "contracts/Escrow.sol";
 
 import {Timestamp} from "contracts/types/Timestamp.sol";
+import {Status} from "contracts/libraries/WithdrawalBatchesQueue.sol";
 
 import "contracts/model/StETHModel.sol";
 import "contracts/model/WstETHAdapted.sol";
@@ -110,17 +111,19 @@ contract StorageSetup is KontrolTest {
         vm.assume(_getRageQuitExtensionDelay(_signallingEscrow) == 0);
         vm.assume(_getRageQuitWithdrawalsTimelock(_signallingEscrow) == 0);
         vm.assume(_getRageQuitTimelockStartedAt(_signallingEscrow) == 0);
+        vm.assume(_getBatchesQueue(_signallingEscrow) == Status.Empty);
     }
 
     function _rageQuitEscrowStorageSetup(IEscrow _rageQuitEscrow, DualGovernance _dualGovernance) internal {
         _escrowStorageSetup(_rageQuitEscrow, _dualGovernance, EscrowState.RageQuitEscrow);
+        vm.assume(_getBatchesQueue(_rageQuitEscrow) != Status.Empty);
     }
 
-    function _getCurrentState(Escrow _escrow) internal view returns (EscrowState) {
+    function _getCurrentState(IEscrow _escrow) internal view returns (EscrowState) {
         return EscrowState(uint8(uint256(vm.load(address(_escrow), 0))));
     }
 
-    function _getLastAssetsLockTimestamp(Escrow _escrow, address _vetoer) internal view returns (uint256) {
+    function _getLastAssetsLockTimestamp(IEscrow _escrow, address _vetoer) internal view returns (uint256) {
         uint256 assetsSlot = 3;
         uint256 vetoerAddressPadded = uint256(uint160(_vetoer));
         bytes32 vetoerAssetsSlot = keccak256(abi.encodePacked(vetoerAddressPadded, assetsSlot));
@@ -140,6 +143,10 @@ contract StorageSetup is KontrolTest {
         return uint40(_loadUInt256(address(_escrow), 9) >> 64);
     }
 
+    function _getBatchesQueue(IEscrow _escrow) internal view returns (Status) {
+        return Status(uint8(_loadUInt256(address(_escrow), 5)));
+    }
+
     function _escrowStorageSetup(IEscrow _escrow, DualGovernance _dualGovernance, EscrowState _currentState) internal {
         kevm.symbolicStorage(address(_escrow));
         // Slot 0
@@ -177,6 +184,12 @@ contract StorageSetup is KontrolTest {
             }
             _storeBytes32(address(_escrow), 2, slot2);
         }
+        // Slot 5
+        {
+            uint256 batchesQueue = kevm.freshUInt(32);
+            vm.assume(batchesQueue < 3);
+            _storeUInt256(address(_escrow), 5, batchesQueue);
+        }
         // Slot 9
         {
             uint32 rageQuitExtensionDelay = uint32(kevm.freshUInt(4));