From 483a7817cdfa012bc599ae2fbf8d48aa4b091972 Mon Sep 17 00:00:00 2001
From: brockelmore <31553173+brockelmore@users.noreply.github.com>
Date: Tue, 14 Dec 2021 17:47:59 -0800
Subject: [PATCH] fix bad state reset (#226)

* fix bad state reset

* fmt

* clippy
---
 forge/src/multi_runner.rs |  7 +++++--
 forge/src/runner.rs       | 19 ++++++++++++++-----
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/forge/src/multi_runner.rs b/forge/src/multi_runner.rs
index 439b95ddcb97..dd026b469064 100644
--- a/forge/src/multi_runner.rs
+++ b/forge/src/multi_runner.rs
@@ -133,10 +133,12 @@ where
     ) -> Result<BTreeMap<String, BTreeMap<String, TestResult>>> {
         // TODO: Convert to iterator, ideally parallel one?
         let contracts = std::mem::take(&mut self.contracts);
+
+        let init_state: S = self.evm.state().clone();
         let results = contracts
             .iter()
             .map(|(name, (abi, address, logs))| {
-                let result = self.run_tests(name, abi, *address, logs, &pattern)?;
+                let result = self.run_tests(name, abi, *address, logs, &pattern, &init_state)?;
                 Ok((name.clone(), result))
             })
             .filter_map(|x: Result<_>| x.ok())
@@ -162,10 +164,11 @@ where
         address: Address,
         init_logs: &[String],
         pattern: &Regex,
+        init_state: &S,
     ) -> Result<BTreeMap<String, TestResult>> {
         let mut runner =
             ContractRunner::new(&mut self.evm, contract, address, self.sender, init_logs);
-        runner.run_tests(pattern, self.fuzzer.as_mut())
+        runner.run_tests(pattern, self.fuzzer.as_mut(), init_state)
     }
 }
 
diff --git a/forge/src/runner.rs b/forge/src/runner.rs
index e24acded946b..305805d1e72b 100644
--- a/forge/src/runner.rs
+++ b/forge/src/runner.rs
@@ -96,6 +96,7 @@ impl<'a, S: Clone, E: Evm<S>> ContractRunner<'a, S, E> {
         &mut self,
         regex: &Regex,
         fuzzer: Option<&mut TestRunner>,
+        init_state: &S,
     ) -> Result<BTreeMap<String, TestResult>> {
         tracing::info!("starting tests");
         let start = Instant::now();
@@ -108,8 +109,6 @@ impl<'a, S: Clone, E: Evm<S>> ContractRunner<'a, S, E> {
             .filter(|func| regex.is_match(&func.name))
             .collect::<Vec<_>>();
 
-        let init_state = self.evm.state().clone();
-
         // run all unit tests
         let unit_tests = test_fns
             .iter()
@@ -284,6 +283,8 @@ mod tests {
                 .deploy(Address::zero(), compiled.bytecode().unwrap().clone(), 0.into())
                 .unwrap();
 
+            let init_state = evm.state().clone();
+
             let mut runner =
                 ContractRunner::new(&mut evm, compiled.abi.as_ref().unwrap(), addr, None, &[]);
 
@@ -291,7 +292,11 @@ mod tests {
             cfg.failure_persistence = None;
             let mut fuzzer = TestRunner::new(cfg);
             let results = runner
-                .run_tests(&Regex::from_str("testGreeting").unwrap(), Some(&mut fuzzer))
+                .run_tests(
+                    &Regex::from_str("testGreeting").unwrap(),
+                    Some(&mut fuzzer),
+                    &init_state,
+                )
                 .unwrap();
             assert!(results["testGreeting()"].success);
             assert!(results["testGreeting(string)"].success);
@@ -311,6 +316,8 @@ mod tests {
                 .deploy(Address::zero(), compiled.bytecode().unwrap().clone(), 0.into())
                 .unwrap();
 
+            let init_state = evm.state().clone();
+
             let mut runner =
                 ContractRunner::new(&mut evm, compiled.abi.as_ref().unwrap(), addr, None, &[]);
 
@@ -318,7 +325,7 @@ mod tests {
             cfg.failure_persistence = None;
             let mut fuzzer = TestRunner::new(cfg);
             let results = runner
-                .run_tests(&Regex::from_str("testFuzz.*").unwrap(), Some(&mut fuzzer))
+                .run_tests(&Regex::from_str("testFuzz.*").unwrap(), Some(&mut fuzzer), &init_state)
                 .unwrap();
             for (_, res) in results {
                 assert!(!res.success);
@@ -422,10 +429,12 @@ mod tests {
         let (addr, _, _, _) =
             evm.deploy(Address::zero(), compiled.bytecode().unwrap().clone(), 0.into()).unwrap();
 
+        let init_state = evm.state().clone();
+
         let mut runner =
             ContractRunner::new(&mut evm, compiled.abi.as_ref().unwrap(), addr, None, &[]);
 
-        let res = runner.run_tests(&".*".parse().unwrap(), None).unwrap();
+        let res = runner.run_tests(&".*".parse().unwrap(), None, &init_state).unwrap();
         assert!(!res.is_empty());
         assert!(res.iter().all(|(_, result)| result.success));
     }