Support assert revert of the given contract at any subcall of the nex…

…t call
foundry-rs · Sep 16, 2024 · 6e3b7ff · 6e3b7ff
1 parent 3336ec0
commit 6e3b7ff
Show file tree

Hide file tree

Showing 4 changed files with 425 additions and 159 deletions.
diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs
@@ -732,7 +732,6 @@ impl Cheatcodes {
                     &expected_revert,
                     outcome.result.result,
                     outcome.result.output.clone(),
-                    None,
                     &self.config.available_artifacts,
                 ) {
                     Ok((address, retdata)) => {
@@ -1234,8 +1233,17 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
             }
         }
 
-        // Handle expected reverts
-        if let Some(expected_revert) = &self.expected_revert {
+        // Handle expected reverts.
+        if let Some(expected_revert) = &mut self.expected_revert {
+            // Record current reverter address before processing the expect revert if call reverted,
+            // expect revert is set with expected reverter address and no actual reverter set yet.
+            if outcome.result.is_revert() &&
+                expected_revert.reverter.is_some() &&
+                expected_revert.reverted_by.is_none()
+            {
+                expected_revert.reverted_by = Some(call.target_address);
+            }
+
             if ecx.journaled_state.depth() <= expected_revert.depth {
                 let needs_processing = match expected_revert.kind {
                     ExpectedRevertKind::Default => !cheatcode_call,
@@ -1254,7 +1262,6 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
                         &expected_revert,
                         outcome.result.result,
                         outcome.result.output.clone(),
-                        Some(call.target_address),
                         &self.config.available_artifacts,
                     ) {
                         Err(error) => {

diff --git a/crates/cheatcodes/src/test/expect.rs b/crates/cheatcodes/src/test/expect.rs
@@ -82,6 +82,8 @@ pub struct ExpectedRevert {
     pub partial_match: bool,
     /// Contract expected to revert next call.
     pub reverter: Option<Address>,
+    /// Actual reverter of the call.
+    pub reverted_by: Option<Address>,
 }
 
 #[derive(Clone, Debug)]
@@ -672,6 +674,7 @@ fn expect_revert(
         },
         partial_match,
         reverter,
+        reverted_by: None,
     });
     Ok(Default::default())
 }
@@ -682,7 +685,6 @@ pub(crate) fn handle_expect_revert(
     expected_revert: &ExpectedRevert,
     status: InstructionResult,
     retdata: Bytes,
-    target_address: Option<Address>,
     known_contracts: &Option<ContractsByArtifact>,
 ) -> Result<(Option<Address>, Bytes)> {
     let success_return = || {
@@ -697,7 +699,7 @@ pub(crate) fn handle_expect_revert(
 
     // If expected reverter address is set then check it matches the actual reverter.
     if let (Some(expected_reverter), Some(actual_reverter)) =
-        (expected_revert.reverter, target_address)
+        (expected_revert.reverter, expected_revert.reverted_by)
     {
         if expected_reverter != actual_reverter {
             return Err(fmt_err!(