feat(test): allow custom txes before unit and fuzz test (#8497)

grandizzy · mattsse · web-flow · commit 4a41367398dd · 2024-07-26T10:26:21.000Z
* feat(test): allow performing txes before unit test

* Changes after review:
- do not unwrap func
- check if `beforeTestSelectors` exists
- move logic in prepare_unit_test fn
- apply same logic to fuzz tests

* Review: Before test is not a test kind

* Changes after review: beforeTestSetup new fn signature

* Remove obsolete struct from test

* Update crates/forge/src/runner.rs

Co-authored-by: Matthias Seitz &lt;matthias.seitz@outlook.de&gt;

* Changes after review: avoid executor clone

* Fix Cow::Borrowed usage

---------

Co-authored-by: Matthias Seitz &lt;matthias.seitz@outlook.de&gt;
diff --git a/crates/common/src/traits.rs b/crates/common/src/traits.rs
@@ -44,6 +44,11 @@ pub trait TestFunctionExt {
         matches!(self.test_function_kind(), TestFunctionKind::UnitTest { .. })
     }
 
+    /// Returns `true` if this function is a `beforeTestSetup` function.
+    fn is_before_test_setup(&self) -> bool {
+        self.tfe_as_str().eq_ignore_ascii_case("beforetestsetup")
+    }
+
     /// Returns `true` if this function is a fuzz test.
     fn is_fuzz_test(&self) -> bool {
         self.test_function_kind().is_fuzz_test()
diff --git a/crates/evm/evm/src/executors/invariant/mod.rs b/crates/evm/evm/src/executors/invariant/mod.rs
@@ -532,6 +532,7 @@ impl<'a> InvariantExecutor<'a> {
     /// targetArtifactSelectors > excludeArtifacts > targetArtifacts
     pub fn select_contract_artifacts(&mut self, invariant_address: Address) -> Result<()> {
         let result = self
+            .executor
             .call_sol_default(invariant_address, &IInvariantTest::targetArtifactSelectorsCall {});
 
         // Insert them into the executor `targeted_abi`.
@@ -542,10 +543,12 @@ impl<'a> InvariantExecutor<'a> {
             self.artifact_filters.targeted.entry(identifier).or_default().extend(selectors);
         }
 
-        let selected =
-            self.call_sol_default(invariant_address, &IInvariantTest::targetArtifactsCall {});
-        let excluded =
-            self.call_sol_default(invariant_address, &IInvariantTest::excludeArtifactsCall {});
+        let selected = self
+            .executor
+            .call_sol_default(invariant_address, &IInvariantTest::targetArtifactsCall {});
+        let excluded = self
+            .executor
+            .call_sol_default(invariant_address, &IInvariantTest::excludeArtifactsCall {});
 
         // Insert `excludeArtifacts` into the executor `excluded_abi`.
         for contract in excluded.excludedArtifacts {
@@ -620,10 +623,14 @@ impl<'a> InvariantExecutor<'a> {
         &self,
         to: Address,
     ) -> Result<(SenderFilters, FuzzRunIdentifiedContracts)> {
-        let targeted_senders =
-            self.call_sol_default(to, &IInvariantTest::targetSendersCall {}).targetedSenders;
-        let mut excluded_senders =
-            self.call_sol_default(to, &IInvariantTest::excludeSendersCall {}).excludedSenders;
+        let targeted_senders = self
+            .executor
+            .call_sol_default(to, &IInvariantTest::targetSendersCall {})
+            .targetedSenders;
+        let mut excluded_senders = self
+            .executor
+            .call_sol_default(to, &IInvariantTest::excludeSendersCall {})
+            .excludedSenders;
         // Extend with default excluded addresses - https://github.com/foundry-rs/foundry/issues/4163
         excluded_senders.extend([
             CHEATCODE_ADDRESS,
@@ -634,10 +641,14 @@ impl<'a> InvariantExecutor<'a> {
         excluded_senders.extend(PRECOMPILES);
         let sender_filters = SenderFilters::new(targeted_senders, excluded_senders);
 
-        let selected =
-            self.call_sol_default(to, &IInvariantTest::targetContractsCall {}).targetedContracts;
-        let excluded =
-            self.call_sol_default(to, &IInvariantTest::excludeContractsCall {}).excludedContracts;
+        let selected = self
+            .executor
+            .call_sol_default(to, &IInvariantTest::targetContractsCall {})
+            .targetedContracts;
+        let excluded = self
+            .executor
+            .call_sol_default(to, &IInvariantTest::excludeContractsCall {})
+            .excludedContracts;
 
         let contracts = self
             .setup_contracts
@@ -678,6 +689,7 @@ impl<'a> InvariantExecutor<'a> {
         targeted_contracts: &mut TargetedContracts,
     ) -> Result<()> {
         let interfaces = self
+            .executor
             .call_sol_default(invariant_address, &IInvariantTest::targetInterfacesCall {})
             .targetedInterfaces;
 
@@ -735,13 +747,15 @@ impl<'a> InvariantExecutor<'a> {
         }
 
         // Collect contract functions marked as target for fuzzing campaign.
-        let selectors = self.call_sol_default(address, &IInvariantTest::targetSelectorsCall {});
+        let selectors =
+            self.executor.call_sol_default(address, &IInvariantTest::targetSelectorsCall {});
         for IInvariantTest::FuzzSelector { addr, selectors } in selectors.targetedSelectors {
             self.add_address_with_functions(addr, &selectors, false, targeted_contracts)?;
         }
 
         // Collect contract functions excluded from fuzzing campaign.
-        let selectors = self.call_sol_default(address, &IInvariantTest::excludeSelectorsCall {});
+        let selectors =
+            self.executor.call_sol_default(address, &IInvariantTest::excludeSelectorsCall {});
         for IInvariantTest::FuzzSelector { addr, selectors } in selectors.excludedSelectors {
             self.add_address_with_functions(addr, &selectors, true, targeted_contracts)?;
         }
@@ -773,17 +787,6 @@ impl<'a> InvariantExecutor<'a> {
         contract.add_selectors(selectors.iter().copied(), should_exclude)?;
         Ok(())
     }
-
-    fn call_sol_default<C: SolCall>(&self, to: Address, args: &C) -> C::Return
-    where
-        C::Return: Default,
-    {
-        self.executor
-            .call_sol(CALLER, to, args, U256::ZERO, None)
-            .map(|c| c.decoded_result)
-            .inspect_err(|e| warn!(target: "forge::test", "failed calling {:?}: {e}", C::SIGNATURE))
-            .unwrap_or_default()
-    }
 }
 
 /// Collects data from call for fuzzing. However, it first verifies that the sender is not an EOA
diff --git a/crates/evm/evm/src/executors/mod.rs b/crates/evm/evm/src/executors/mod.rs
@@ -50,6 +50,9 @@ sol! {
     interface ITest {
         function setUp() external;
         function failed() external view returns (bool failed);
+
+        #[derive(Default)]
+        function beforeTestSetup(bytes4 testSelector) public view returns (bytes[] memory beforeTestCalldata);
     }
 }
 
@@ -602,6 +605,16 @@ impl Executor {
 
         EnvWithHandlerCfg::new_with_spec_id(Box::new(env), self.spec_id())
     }
+
+    pub fn call_sol_default<C: SolCall>(&self, to: Address, args: &C) -> C::Return
+    where
+        C::Return: Default,
+    {
+        self.call_sol(CALLER, to, args, U256::ZERO, None)
+            .map(|c| c.decoded_result)
+            .inspect_err(|e| warn!(target: "forge::test", "failed calling {:?}: {e}", C::SIGNATURE))
+            .unwrap_or_default()
+    }
 }
 
 /// Represents the context after an execution error occurred.
diff --git a/crates/forge/src/result.rs b/crates/forge/src/result.rs
@@ -449,9 +449,9 @@ impl TestResult {
     }
 
     /// Returns the failed result with reason for single test.
-    pub fn single_fail(mut self, err: EvmError) -> Self {
+    pub fn single_fail(mut self, reason: Option<String>) -> Self {
         self.status = TestStatus::Failure;
-        self.reason = Some(err.to_string());
+        self.reason = reason;
         self
     }
 
@@ -579,6 +579,14 @@ impl TestResult {
         format!("{self} {name} {}", self.kind.report())
     }
 
+    /// Function to merge logs, addresses, traces and coverage from a call result into test result.
+    pub fn merge_call_result(&mut self, call_result: &RawCallResult) {
+        self.logs.extend(call_result.logs.clone());
+        self.labeled_addresses.extend(call_result.labels.clone());
+        self.traces.extend(call_result.traces.clone().map(|traces| (TraceKind::Execution, traces)));
+        self.merge_coverages(call_result.coverage.clone());
+    }
+
     /// Function to merge given coverage in current test result coverage.
     pub fn merge_coverages(&mut self, other_coverage: Option<HitMaps>) {
         let old_coverage = std::mem::take(&mut self.coverage);
diff --git a/crates/forge/src/runner.rs b/crates/forge/src/runner.rs
@@ -24,7 +24,7 @@ use foundry_evm::{
         invariant::{
             check_sequence, replay_error, replay_run, InvariantExecutor, InvariantFuzzError,
         },
-        CallResult, EvmError, ExecutionErr, Executor, RawCallResult,
+        CallResult, EvmError, ExecutionErr, Executor, ITest, RawCallResult,
     },
     fuzz::{
         fixture_name,
@@ -36,6 +36,7 @@ use foundry_evm::{
 use proptest::test_runner::TestRunner;
 use rayon::prelude::*;
 use std::{
+    borrow::Cow,
     cmp::min,
     collections::{BTreeMap, HashMap},
     time::Instant,
@@ -270,6 +271,7 @@ impl<'a> ContractRunner<'a> {
                 ));
             }
         }
+
         // There are multiple setUp function, so we return a single test result for `setUp`
         if setup_fns.len() > 1 {
             return SuiteResult::new(
@@ -412,21 +414,26 @@ impl<'a> ContractRunner<'a> {
 
     /// Runs a single unit test.
     ///
-    /// Calls the given functions and returns the `TestResult`.
+    /// Applies before test txes (if any), runs current test and returns the `TestResult`.
     ///
-    /// State modifications are not committed to the evm database but discarded after the call,
-    /// similar to `eth_call`.
+    /// Before test txes are applied in order and state modifications committed to the EVM database
+    /// (therefore the unit test call will be made on modified state).
+    /// State modifications of before test txes and unit test function call are discarded after
+    /// test ends, similar to `eth_call`.
     pub fn run_unit_test(
         &self,
         func: &Function,
         should_fail: bool,
         setup: TestSetup,
     ) -> TestResult {
-        let address = setup.address;
-        let test_result = TestResult::new(setup);
+        // Prepare unit test execution.
+        let (executor, test_result, address) = match self.prepare_test(func, setup) {
+            Ok(res) => res,
+            Err(res) => return res,
+        };
 
-        // Run unit test
-        let (mut raw_call_result, reason) = match self.executor.call(
+        // Run current unit test.
+        let (mut raw_call_result, reason) = match executor.call(
             self.sender,
             address,
             func,
@@ -437,11 +444,10 @@ impl<'a> ContractRunner<'a> {
             Ok(res) => (res.raw, None),
             Err(EvmError::Execution(err)) => (err.raw, Some(err.reason)),
             Err(EvmError::SkipError) => return test_result.single_skip(),
-            Err(err) => return test_result.single_fail(err),
+            Err(err) => return test_result.single_fail(Some(err.to_string())),
         };
 
-        let success =
-            self.executor.is_raw_call_mut_success(address, &mut raw_call_result, should_fail);
+        let success = executor.is_raw_call_mut_success(address, &mut raw_call_result, should_fail);
         test_result.single_result(success, reason, raw_call_result)
     }
 
@@ -618,6 +624,15 @@ impl<'a> ContractRunner<'a> {
         )
     }
 
+    /// Runs a fuzzed test.
+    ///
+    /// Applies the before test txes (if any), fuzzes the current function and returns the
+    /// `TestResult`.
+    ///
+    /// Before test txes are applied in order and state modifications committed to the EVM database
+    /// (therefore the fuzz test will use the modified state).
+    /// State modifications of before test txes and fuzz test are discarded after test ends,
+    /// similar to `eth_call`.
     pub fn run_fuzz_test(
         &self,
         func: &Function,
@@ -626,14 +641,18 @@ impl<'a> ContractRunner<'a> {
         setup: TestSetup,
         fuzz_config: FuzzConfig,
     ) -> TestResult {
-        let address = setup.address;
+        let progress = start_fuzz_progress(self.progress, self.name, &func.name, fuzz_config.runs);
+
+        // Prepare fuzz test execution.
         let fuzz_fixtures = setup.fuzz_fixtures.clone();
-        let test_result = TestResult::new(setup);
+        let (executor, test_result, address) = match self.prepare_test(func, setup) {
+            Ok(res) => res,
+            Err(res) => return res,
+        };
 
-        // Run fuzz test
-        let progress = start_fuzz_progress(self.progress, self.name, &func.name, fuzz_config.runs);
+        // Run fuzz test.
         let fuzzed_executor =
-            FuzzedExecutor::new(self.executor.clone(), runner, self.sender, fuzz_config);
+            FuzzedExecutor::new(executor.into_owned(), runner, self.sender, fuzz_config);
         let result = fuzzed_executor.fuzz(
             func,
             &fuzz_fixtures,
@@ -650,4 +669,51 @@ impl<'a> ContractRunner<'a> {
         }
         test_result.fuzz_result(result)
     }
+
+    /// Prepares single unit test and fuzz test execution:
+    /// - set up the test result and executor
+    /// - check if before test txes are configured and apply them in order
+    ///
+    /// Before test txes are arrays of arbitrary calldata obtained by calling the `beforeTest`
+    /// function with test selector as a parameter.
+    ///
+    /// Unit tests within same contract (or even current test) are valid options for before test tx
+    /// configuration. Test execution stops if any of before test txes fails.
+    fn prepare_test(
+        &self,
+        func: &Function,
+        setup: TestSetup,
+    ) -> Result<(Cow<'_, Executor>, TestResult, Address), TestResult> {
+        let address = setup.address;
+        let mut executor = Cow::Borrowed(&self.executor);
+        let mut test_result = TestResult::new(setup);
+
+        // Apply before test configured functions (if any).
+        if self.contract.abi.functions().filter(|func| func.name.is_before_test_setup()).count() ==
+            1
+        {
+            for calldata in executor
+                .call_sol_default(
+                    address,
+                    &ITest::beforeTestSetupCall { testSelector: func.selector() },
+                )
+                .beforeTestCalldata
+            {
+                // Apply before test configured calldata.
+                match executor.to_mut().transact_raw(self.sender, address, calldata, U256::ZERO) {
+                    Ok(call_result) => {
+                        // Merge tx result traces in unit test result.
+                        test_result.merge_call_result(&call_result);
+
+                        // To continue unit test execution the call should not revert.
+                        if call_result.reverted {
+                            return Err(test_result.single_fail(None))
+                        }
+                    }
+                    Err(_) => return Err(test_result.single_fail(None)),
+                }
+            }
+        }
+        Ok((executor, test_result, address))
+    }
 }
diff --git a/crates/forge/tests/it/repros.rs b/crates/forge/tests/it/repros.rs
@@ -364,3 +364,6 @@ test_repro!(8168);
 
 // https://github.com/foundry-rs/foundry/issues/8383
 test_repro!(8383);
+
+// https://github.com/foundry-rs/foundry/issues/1543
+test_repro!(1543);
diff --git a/testdata/default/repros/Issue1543.t.sol b/testdata/default/repros/Issue1543.t.sol

Original file line number	Diff line number	Diff line change
`@@ -50,6 +50,9 @@ sol! {`
`50`	`50`	`interface ITest {`
`51`	`51`	`function setUp() external;`
`52`	`52`	`function failed() external view returns (bool failed);`
	`53`	`+`
	`54`	`+ #[derive(Default)]`
	`55`	`+ function beforeTestSetup(bytes4 testSelector) public view returns (bytes[] memory beforeTestCalldata);`
`53`	`56`	`}`
`54`	`57`	`}`
`55`	`58`
`@@ -602,6 +605,16 @@ impl Executor {`
`602`	`605`
`603`	`606`	`EnvWithHandlerCfg::new_with_spec_id(Box::new(env), self.spec_id())`
`604`	`607`	`}`
	`608`	`+`
	`609`	`+ pub fn call_sol_default<C: SolCall>(&self, to: Address, args: &C) -> C::Return`
	`610`	`+ where`
	`611`	`+ C::Return: Default,`
	`612`	`+ {`
	`613`	`+ self.call_sol(CALLER, to, args, U256::ZERO, None)`
	`614`	`+ .map(\|c\| c.decoded_result)`
	`615`	`+ .inspect_err(\|e\| warn!(target: "forge::test", "failed calling {:?}: {e}", C::SIGNATURE))`
	`616`	`+ .unwrap_or_default()`
	`617`	`+ }`
`605`	`618`	`}`
`606`	`619`
`607`	`620`	`/// Represents the context after an execution error occurred.`