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); + } +}