From 65b1f01cef3de014a2d6646ebe0ca0df202cc9e6 Mon Sep 17 00:00:00 2001
From: Thomas Clement <88335455+tom2o17@users.noreply.github.com>
Date: Mon, 5 Feb 2024 15:31:39 -0500
Subject: [PATCH] Fix useMultipleSetup (#17)

Reverting to the first snapshot in the list, also implicitly resets the array of snapShotIds to Only have one member
---
 src/FraxTest.sol             |  3 ++-
 test/TestMultipleSetup.t.sol | 31 +++++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)
 create mode 100644 test/TestMultipleSetup.t.sol

diff --git a/src/FraxTest.sol b/src/FraxTest.sol
index c159f8f..7bd0132 100644
--- a/src/FraxTest.sol
+++ b/src/FraxTest.sol
@@ -12,16 +12,17 @@ abstract contract FraxTest is VmHelper, Test {
     modifier useMultipleSetupFunctions() {
         if (snapShotIds.length == 0) _;
         for (uint256 i = 0; i < snapShotIds.length; i++) {
+            uint256 _originalSnapshotId = vm.snapshot();
             if (!vm.revertTo(snapShotIds[i])) {
                 revert VmDidNotRevert(snapShotIds[i]);
             }
             _;
             vm.clearMockedCalls();
+            vm.revertTo(_originalSnapshotId);
         }
     }
 
     function addSetupFunctions(function()[] memory _setupFunctions) internal {
-        uint256 _originalSnapshotId = vm.snapshot();
         for (uint256 i = 0; i < _setupFunctions.length; i++) {
             _setupFunctions[i]();
             snapShotIds.push(vm.snapshot());
diff --git a/test/TestMultipleSetup.t.sol b/test/TestMultipleSetup.t.sol
new file mode 100644
index 0000000..c9f6a23
--- /dev/null
+++ b/test/TestMultipleSetup.t.sol
@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: ISC
+pragma solidity >=0.8.0;
+
+import "../src/FraxTest.sol";
+
+contract TestMultipleSetup is FraxTest {
+    uint256 value;
+
+    function initializeValueOne() public {
+        value = 25;
+    }
+
+    function initializeValueTwo() public {
+        value = 5;
+    }
+
+    function initializeValueThree() public {
+        value = 10;
+    }
+
+    function setUp() public {
+        setupFunctions.push(initializeValueOne);
+        setupFunctions.push(initializeValueTwo);
+        setupFunctions.push(initializeValueThree);
+        addSetupFunctions(setupFunctions);
+    }
+
+    function testFailAssertValue() public useMultipleSetupFunctions {
+        assertEq(value, 5);
+    }
+}