allow combining expectRevert with assumeNoRevert

foundry-rs · Jan 3, 2025 · 0ea0fb2 · 0ea0fb2
1 parent a938ad1
commit 0ea0fb2
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 24 deletions.
diff --git a/crates/cheatcodes/src/test/assume.rs b/crates/cheatcodes/src/test/assume.rs
@@ -6,7 +6,6 @@ use spec::Vm::{
 };
 use std::fmt::Debug;
 
-pub const ASSUME_EXPECT_REJECT_MAGIC: &str = "Cannot combine an assumeNoRevert with expectRevert";
 pub const ASSUME_REJECT_MAGIC: &str =
     "Cannot combine a generic assumeNoRevert with specific assumeNoRevert reasons";
 
@@ -90,7 +89,6 @@ fn assume_no_revert(
     depth: u64,
     parameters: Vec<AcceptableRevertParameters>,
 ) -> Result {
-    ensure!(state.expected_revert.is_none(), ASSUME_EXPECT_REJECT_MAGIC);
     ensure!(state.assume_no_revert.is_none(), ASSUME_REJECT_MAGIC);
 
     state.assume_no_revert = Some(AssumeNoRevert { depth, reasons: parameters, reverted_by: None });

diff --git a/crates/cheatcodes/src/test/expect.rs b/crates/cheatcodes/src/test/expect.rs
@@ -900,7 +900,6 @@ fn expect_revert(
         state.expected_revert.is_none(),
         "you must call another function prior to expecting a second revert"
     );
-    ensure!(state.assume_no_revert.is_none(), "Cannot expect a revert when using assumeNoRevert");
     state.expected_revert = Some(ExpectedRevert {
         reason: reason.map(<[_]>::to_vec),
         depth,

diff --git a/crates/forge/tests/cli/test_cmd.rs b/crates/forge/tests/cli/test_cmd.rs
@@ -2480,6 +2480,7 @@ interface Vm {
     function assumeNoRevert() external pure;
     function assumeNoRevert(PotentialRevert calldata revertData) external pure;
     function assumeNoRevert(PotentialRevert[] calldata revertData) external pure;
+    function expectRevert(bytes4 revertData, uint64 count) external;
 }
 
 contract ReverterB {
@@ -2503,6 +2504,7 @@ contract Reverter {
     error MyRevert();
     error RevertWithData(uint256 x);
     error UnusedError();
+    error ExpectedRevertCountZero();
 
     ReverterB public immutable subReverter;
 
@@ -2536,6 +2538,15 @@ contract Reverter {
         }
         return true;
     }
+
+    function revertIf2Or3ExpectedRevertZero(uint256 x) public pure returns (bool) {
+        if (x == 2) {
+            revert ExpectedRevertCountZero();
+        } else if (x == 3) {
+            revert MyRevert();
+        }
+        return true;
+    }
 }
 
 contract ReverterTest is Test {
@@ -2592,13 +2603,31 @@ contract ReverterTest is Test {
         _vm.assumeNoRevert();
         reverter.twoPossibleReverts(2);
     }
-
-    /// @dev Test that calling `expectRevert` after `assumeNoRevert` results in an error
-    function testAssumeThenExpect_fails(uint256) public {
-        _vm.assumeNoRevert(Vm.PotentialRevert({revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), partialMatch: false, reverter: address(0)}));
-        _vm.expectRevert();
-        reverter.revertIf2(1);
+    
+    function testAssumeThenExpectCountZeroFails(uint256 x) public {
+        _vm.assumeNoRevert(
+            Vm.PotentialRevert({
+                revertData: abi.encodeWithSelector(Reverter.MyRevert.selector),
+                partialMatch: false,
+                reverter: address(0)
+            })
+        );
+        _vm.expectRevert(Reverter.ExpectedRevertCountZero.selector, 0);
+        reverter.revertIf2Or3ExpectedRevertZero(x);
+    }
+
+    function testExpectCountZeroThenAssumeFails(uint256 x) public {
+        _vm.expectRevert(Reverter.ExpectedRevertCountZero.selector, 0);
+        _vm.assumeNoRevert(
+            Vm.PotentialRevert({
+                revertData: abi.encodeWithSelector(Reverter.MyRevert.selector),
+                partialMatch: false,
+                reverter: address(0)
+            })
+        );
+        reverter.revertIf2Or3ExpectedRevertZero(x);
     }
+
 }"#,
     )
     .unwrap();
@@ -2607,11 +2636,12 @@ contract ReverterTest is Test {
 [SOLC_VERSION] [ELAPSED]
 Compiler run successful!
 
-Ran 7 tests for src/AssumeNoRevertTest.t.sol:ReverterTest
-[FAIL: vm.expectRevert: Cannot expect a revert when using assumeNoRevert; counterexample: [..]] testAssumeThenExpect_fails(uint256) (runs: [..], [AVG_GAS])
+Ran 8 tests for src/AssumeNoRevertTest.t.sol:ReverterTest
+[FAIL: expected 0 reverts with reason: 0x92fa317b, but got one; counterexample: [..]] testAssumeThenExpectCountZeroFails(uint256) (runs: [..], [AVG_GAS])
 [FAIL: MyRevert(); counterexample: calldata=[..]] testAssumeWithReverter_fails(uint256) (runs: [..], [AVG_GAS])
 [FAIL: RevertWithData(2); counterexample: [..]] testAssume_wrongData_fails(uint256) (runs: [..], [AVG_GAS])
 [FAIL: MyRevert(); counterexample: [..]] testAssume_wrongSelector_fails(uint256) (runs: [..], [AVG_GAS])
+[FAIL: expected 0 reverts with reason: 0x92fa317b, but got one; counterexample: [..]] testExpectCountZeroThenAssumeFails(uint256) (runs: [..], [AVG_GAS])
 [FAIL: MyRevert(); counterexample: [..]] testMultipleAssumesClearAfterCall_fails(uint256) (runs: 0, [AVG_GAS])
 [FAIL: RevertWithData(3); counterexample: [..]] testMultipleAssumes_OneWrong_fails(uint256) (runs: [..], [AVG_GAS])
 [FAIL: vm.assumeNoRevert: Cannot combine a generic assumeNoRevert with specific assumeNoRevert reasons; counterexample: [..]] testMultipleAssumes_ThrowOnGenericNoRevert_AfterSpecific_fails(bytes4) (runs: [..], [AVG_GAS])

diff --git a/testdata/default/cheats/AssumeNoRevert.t.sol b/testdata/default/cheats/AssumeNoRevert.t.sol
@@ -150,17 +150,4 @@ contract ReverterTest is Test {
         _vm.assumeNoRevert(revertData);
         reverter.twoPossibleReverts(x);
     }
-
-    /// @dev Test that calling `assumeNoRevert` after `expectRevert` results in an error
-    function testExpectThenAssumeFails() public {
-        _vm._expectCheatcodeRevert();
-        _vm.assumeNoRevert(
-            Vm.PotentialRevert({
-                revertData: abi.encodeWithSelector(Reverter.MyRevert.selector),
-                partialMatch: false,
-                reverter: address(0)
-            })
-        );
-        reverter.revertIf2(1);
-    }
 }