fix(forge) expectRevert for cheatcodes

crates/cheatcodes/src/inspector.rs

@@ -919,6 +919,30 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
         status: InstructionResult,
         retdata: Bytes,
     ) -> (InstructionResult, Gas, Bytes) {
+        // Handle expected reverts
+        if let Some(expected_revert) = &self.expected_revert {
+            if data.journaled_state.depth() <= expected_revert.depth {
+                if expected_revert.pending {
+                    let expected_revert = self.expected_revert.as_mut().unwrap();
+                    expected_revert.pending = false;
+                } else {
+                    let expected_revert = std::mem::take(&mut self.expected_revert).unwrap();
+                    return match expect::handle_expect_revert(
+                        false,
+                        expected_revert.reason.as_deref(),
+                        status,
+                        retdata,
+                    ) {
+                        Err(error) => {
+                            trace!(expected=?expected_revert, ?error, ?status, "Expected revert mismatch");
+                            (InstructionResult::Revert, remaining_gas, error.abi_encode().into())
+                        }
+                        Ok((_, retdata)) => (InstructionResult::Return, remaining_gas, retdata),
+                    };
+                }
+            }
+        }
+
         if call.contract == CHEATCODE_ADDRESS || call.contract == HARDHAT_CONSOLE_ADDRESS {
             return (status, remaining_gas, retdata);
         }
@@ -955,25 +979,6 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
             }
         }
 
-        // Handle expected reverts
-        if let Some(expected_revert) = &self.expected_revert {
-            if data.journaled_state.depth() <= expected_revert.depth {
-                let expected_revert = std::mem::take(&mut self.expected_revert).unwrap();
-                return match expect::handle_expect_revert(
-                    false,
-                    expected_revert.reason.as_deref(),
-                    status,
-                    retdata,
-                ) {
-                    Err(error) => {
-                        trace!(expected=?expected_revert, ?error, ?status, "Expected revert mismatch");
-                        (InstructionResult::Revert, remaining_gas, error.abi_encode().into())
-                    }
-                    Ok((_, retdata)) => (InstructionResult::Return, remaining_gas, retdata),
-                };
-            }
-        }
-
         // If `startStateDiffRecording` has been called, update the `reverted` status of the
         // previous call depth's recorded accesses, if any
         if let Some(recorded_account_diffs_stack) = &mut self.recorded_account_diffs_stack {

crates/cheatcodes/src/test/expect.rs

@@ -59,6 +59,9 @@ pub struct ExpectedRevert {
     pub reason: Option<Vec<u8>>,
     /// The depth at which the revert is expected
     pub depth: u64,
+    /// Flag which is being switched once we exit `expectRevert` call
+    /// Needed to not fail due to `expectRevert` not reverting
+    pub pending: bool,
 }
 
 #[derive(Clone, Debug)]
@@ -435,7 +438,8 @@ fn expect_revert(state: &mut Cheatcodes, reason: Option<&[u8]>, depth: u64) -> R
         state.expected_revert.is_none(),
         "you must call another function prior to expecting a second revert"
     );
-    state.expected_revert = Some(ExpectedRevert { reason: reason.map(<[_]>::to_vec), depth });
+    state.expected_revert =
+        Some(ExpectedRevert { reason: reason.map(<[_]>::to_vec), depth, pending: true });
     Ok(Default::default())
 }
 

testdata/cheats/Json.t.sol

@@ -155,19 +155,27 @@ contract ParseJsonTest is DSTest {
     function test_parseJsonKeys() public {
         string memory jsonString =
             '{"some_key_to_value": "some_value", "some_key_to_array": [1, 2, 3], "some_key_to_object": {"key1": "value1", "key2": 2}}';
+
         string[] memory keys = vm.parseJsonKeys(jsonString, "$");
-        assertEq(abi.encode(keys), abi.encode(["some_key_to_value", "some_key_to_array", "some_key_to_object"]));
+        string[] memory expected = new string[](3);
+        expected[0] = "some_key_to_value";
+        expected[1] = "some_key_to_array";
+        expected[2] = "some_key_to_object";
+        assertEq(abi.encode(keys), abi.encode(expected));
 
         keys = vm.parseJsonKeys(jsonString, ".some_key_to_object");
-        assertEq(abi.encode(keys), abi.encode(["key1", "key2"]));
+        expected = new string[](2);
+        expected[0] = "key1";
+        expected[1] = "key2";
+        assertEq(abi.encode(keys), abi.encode(expected));
 
         vm.expectRevert("JSON value at \".some_key_to_array\" is not an object");
         vm.parseJsonKeys(jsonString, ".some_key_to_array");
 
         vm.expectRevert("JSON value at \".some_key_to_value\" is not an object");
         vm.parseJsonKeys(jsonString, ".some_key_to_value");
 
-        vm.expectRevert("JSON value at \".*\" is not an object");
+        vm.expectRevert("key \".*\" must return exactly one JSON object");
         vm.parseJsonKeys(jsonString, ".*");
     }
 }