forge/
multi_runner.rs

Help
1//! Forge test runner for multiple contracts.
2
3use crate::{
4    progress::TestsProgress, result::SuiteResult, runner::LIBRARY_DEPLOYER, ContractRunner,
5    TestFilter,
6};
7use alloy_json_abi::{Function, JsonAbi};
8use alloy_primitives::{Address, Bytes, U256};
9use eyre::Result;
10use foundry_common::{get_contract_name, shell::verbosity, ContractsByArtifact, TestFunctionExt};
11use foundry_compilers::{
12    artifacts::{Contract, Libraries},
13    compilers::Compiler,
14    Artifact, ArtifactId, ProjectCompileOutput,
15};
16use foundry_config::{Config, InlineConfig};
17use foundry_evm::{
18    backend::Backend,
19    decode::RevertDecoder,
20    executors::{Executor, ExecutorBuilder},
21    fork::CreateFork,
22    inspectors::CheatsConfig,
23    opts::EvmOpts,
24    traces::{InternalTraceMode, TraceMode},
25    Env,
26};
27use foundry_linking::{LinkOutput, Linker};
28use rayon::prelude::*;
29use revm::primitives::hardfork::SpecId;
30use std::{
31    borrow::Borrow,
32    collections::BTreeMap,
33    fmt::Debug,
34    path::Path,
35    sync::{mpsc, Arc},
36    time::Instant,
37};
38
39#[derive(Debug, Clone)]
40pub struct TestContract {
41    pub abi: JsonAbi,
42    pub bytecode: Bytes,
43}
44
45pub type DeployableContracts = BTreeMap<ArtifactId, TestContract>;
46
47/// A multi contract runner receives a set of contracts deployed in an EVM instance and proceeds
48/// to run all test functions in these contracts.
49pub struct MultiContractRunner {
50    /// Mapping of contract name to JsonAbi, creation bytecode and library bytecode which
51    /// needs to be deployed & linked against
52    pub contracts: DeployableContracts,
53    /// Known contracts linked with computed library addresses.
54    pub known_contracts: ContractsByArtifact,
55    /// Revert decoder. Contains all known errors and their selectors.
56    pub revert_decoder: RevertDecoder,
57    /// Libraries to deploy.
58    pub libs_to_deploy: Vec<Bytes>,
59    /// Library addresses used to link contracts.
60    pub libraries: Libraries,
61
62    /// The fork to use at launch
63    pub fork: Option<CreateFork>,
64
65    /// The base configuration for the test runner.
66    pub tcfg: TestRunnerConfig,
67}
68
69impl std::ops::Deref for MultiContractRunner {
70    type Target = TestRunnerConfig;
71
72    fn deref(&self) -> &Self::Target {
73        &self.tcfg
74    }
75}
76
77impl std::ops::DerefMut for MultiContractRunner {
78    fn deref_mut(&mut self) -> &mut Self::Target {
79        &mut self.tcfg
80    }
81}
82
83impl MultiContractRunner {
84    /// Returns an iterator over all contracts that match the filter.
85    pub fn matching_contracts<'a: 'b, 'b>(
86        &'a self,
87        filter: &'b dyn TestFilter,
88    ) -> impl Iterator<Item = (&'a ArtifactId, &'a TestContract)> + 'b {
89        self.contracts.iter().filter(|&(id, c)| matches_contract(id, &c.abi, filter))
90    }
91
92    /// Returns an iterator over all test functions that match the filter.
93    pub fn matching_test_functions<'a: 'b, 'b>(
94        &'a self,
95        filter: &'b dyn TestFilter,
96    ) -> impl Iterator<Item = &'a Function> + 'b {
97        self.matching_contracts(filter)
98            .flat_map(|(_, c)| c.abi.functions())
99            .filter(|func| is_matching_test(func, filter))
100    }
101
102    /// Returns an iterator over all test functions in contracts that match the filter.
103    pub fn all_test_functions<'a: 'b, 'b>(
104        &'a self,
105        filter: &'b dyn TestFilter,
106    ) -> impl Iterator<Item = &'a Function> + 'b {
107        self.contracts
108            .iter()
109            .filter(|(id, _)| filter.matches_path(&id.source) && filter.matches_contract(&id.name))
110            .flat_map(|(_, c)| c.abi.functions())
111            .filter(|func| func.is_any_test())
112    }
113
114    /// Returns all matching tests grouped by contract grouped by file (file -> (contract -> tests))
115    pub fn list(&self, filter: &dyn TestFilter) -> BTreeMap<String, BTreeMap<String, Vec<String>>> {
116        self.matching_contracts(filter)
117            .map(|(id, c)| {
118                let source = id.source.as_path().display().to_string();
119                let name = id.name.clone();
120                let tests = c
121                    .abi
122                    .functions()
123                    .filter(|func| is_matching_test(func, filter))
124                    .map(|func| func.name.clone())
125                    .collect::<Vec<_>>();
126                (source, name, tests)
127            })
128            .fold(BTreeMap::new(), |mut acc, (source, name, tests)| {
129                acc.entry(source).or_default().insert(name, tests);
130                acc
131            })
132    }
133
134    /// Executes _all_ tests that match the given `filter`.
135    ///
136    /// The same as [`test`](Self::test), but returns the results instead of streaming them.
137    ///
138    /// Note that this method returns only when all tests have been executed.
139    pub fn test_collect(
140        &mut self,
141        filter: &dyn TestFilter,
142    ) -> Result<BTreeMap<String, SuiteResult>> {
143        Ok(self.test_iter(filter)?.collect())
144    }
145
146    /// Executes _all_ tests that match the given `filter`.
147    ///
148    /// The same as [`test`](Self::test), but returns the results instead of streaming them.
149    ///
150    /// Note that this method returns only when all tests have been executed.
151    pub fn test_iter(
152        &mut self,
153        filter: &dyn TestFilter,
154    ) -> Result<impl Iterator<Item = (String, SuiteResult)>> {
155        let (tx, rx) = mpsc::channel();
156        self.test(filter, tx, false)?;
157        Ok(rx.into_iter())
158    }
159
160    /// Executes _all_ tests that match the given `filter`.
161    ///
162    /// This will create the runtime based on the configured `evm` ops and create the `Backend`
163    /// before executing all contracts and their tests in _parallel_.
164    ///
165    /// Each Executor gets its own instance of the `Backend`.
166    pub fn test(
167        &mut self,
168        filter: &dyn TestFilter,
169        tx: mpsc::Sender<(String, SuiteResult)>,
170        show_progress: bool,
171    ) -> Result<()> {
172        let tokio_handle = tokio::runtime::Handle::current();
173        trace!("running all tests");
174
175        // The DB backend that serves all the data.
176        let db = Backend::spawn(self.fork.take())?;
177
178        let find_timer = Instant::now();
179        let contracts = self.matching_contracts(filter).collect::<Vec<_>>();
180        let find_time = find_timer.elapsed();
181        debug!(
182            "Found {} test contracts out of {} in {:?}",
183            contracts.len(),
184            self.contracts.len(),
185            find_time,
186        );
187
188        if show_progress {
189            let tests_progress = TestsProgress::new(contracts.len(), rayon::current_num_threads());
190            // Collect test suite results to stream at the end of test run.
191            let results: Vec<(String, SuiteResult)> = contracts
192                .par_iter()
193                .map(|&(id, contract)| {
194                    let _guard = tokio_handle.enter();
195                    tests_progress.inner.lock().start_suite_progress(&id.identifier());
196
197                    let result = self.run_test_suite(
198                        id,
199                        contract,
200                        &db,
201                        filter,
202                        &tokio_handle,
203                        Some(&tests_progress),
204                    );
205
206                    tests_progress
207                        .inner
208                        .lock()
209                        .end_suite_progress(&id.identifier(), result.summary());
210
211                    (id.identifier(), result)
212                })
213                .collect();
214
215            tests_progress.inner.lock().clear();
216
217            results.iter().for_each(|result| {
218                let _ = tx.send(result.to_owned());
219            });
220        } else {
221            contracts.par_iter().for_each(|&(id, contract)| {
222                let _guard = tokio_handle.enter();
223                let result = self.run_test_suite(id, contract, &db, filter, &tokio_handle, None);
224                let _ = tx.send((id.identifier(), result));
225            })
226        }
227
228        Ok(())
229    }
230
231    fn run_test_suite(
232        &self,
233        artifact_id: &ArtifactId,
234        contract: &TestContract,
235        db: &Backend,
236        filter: &dyn TestFilter,
237        tokio_handle: &tokio::runtime::Handle,
238        progress: Option<&TestsProgress>,
239    ) -> SuiteResult {
240        let identifier = artifact_id.identifier();
241        let mut span_name = identifier.as_str();
242
243        if !enabled!(tracing::Level::TRACE) {
244            span_name = get_contract_name(&identifier);
245        }
246        let span = debug_span!("suite", name = %span_name);
247        let span_local = span.clone();
248        let _guard = span_local.enter();
249
250        debug!("start executing all tests in contract");
251
252        let executor = self.tcfg.executor(self.known_contracts.clone(), artifact_id, db.clone());
253        let runner = ContractRunner::new(
254            &identifier,
255            contract,
256            executor,
257            progress,
258            tokio_handle,
259            span,
260            self,
261        );
262        let r = runner.run_tests(filter);
263
264        debug!(duration=?r.duration, "executed all tests in contract");
265
266        r
267    }
268}
269
270/// Configuration for the test runner.
271///
272/// This is modified after instantiation through inline config.
273#[derive(Clone)]
274pub struct TestRunnerConfig {
275    /// Project config.
276    pub config: Arc<Config>,
277    /// Inline configuration.
278    pub inline_config: Arc<InlineConfig>,
279
280    /// EVM configuration.
281    pub evm_opts: EvmOpts,
282    /// EVM environment.
283    pub env: Env,
284    /// EVM version.
285    pub spec_id: SpecId,
286    /// The address which will be used to deploy the initial contracts and send all transactions.
287    pub sender: Address,
288
289    /// Whether to collect coverage info
290    pub coverage: bool,
291    /// Whether to collect debug info
292    pub debug: bool,
293    /// Whether to enable steps tracking in the tracer.
294    pub decode_internal: InternalTraceMode,
295    /// Whether to enable call isolation.
296    pub isolation: bool,
297    /// Whether to enable Odyssey features.
298    pub odyssey: bool,
299}
300
301impl TestRunnerConfig {
302    /// Reconfigures all fields using the given `config`.
303    /// This is for example used to override the configuration with inline config.
304    pub fn reconfigure_with(&mut self, config: Arc<Config>) {
305        debug_assert!(!Arc::ptr_eq(&self.config, &config));
306
307        self.spec_id = config.evm_spec_id();
308        self.sender = config.sender;
309        self.odyssey = config.odyssey;
310        self.isolation = config.isolate;
311
312        // Specific to Forge, not present in config.
313        // TODO: self.evm_opts
314        // TODO: self.env
315        // self.coverage = N/A;
316        // self.debug = N/A;
317        // self.decode_internal = N/A;
318
319        self.config = config;
320    }
321
322    /// Configures the given executor with this configuration.
323    pub fn configure_executor(&self, executor: &mut Executor) {
324        // TODO: See above
325
326        let inspector = executor.inspector_mut();
327        // inspector.set_env(&self.env);
328        if let Some(cheatcodes) = inspector.cheatcodes.as_mut() {
329            cheatcodes.config =
330                Arc::new(cheatcodes.config.clone_with(&self.config, self.evm_opts.clone()));
331        }
332        inspector.tracing(self.trace_mode());
333        inspector.collect_coverage(self.coverage);
334        inspector.enable_isolation(self.isolation);
335        inspector.odyssey(self.odyssey);
336        // inspector.set_create2_deployer(self.evm_opts.create2_deployer);
337
338        // executor.env_mut().clone_from(&self.env);
339        executor.set_spec_id(self.spec_id);
340        // executor.set_gas_limit(self.evm_opts.gas_limit());
341        executor.set_legacy_assertions(self.config.legacy_assertions);
342    }
343
344    /// Creates a new executor with this configuration.
345    pub fn executor(
346        &self,
347        known_contracts: ContractsByArtifact,
348        artifact_id: &ArtifactId,
349        db: Backend,
350    ) -> Executor {
351        let cheats_config = Arc::new(CheatsConfig::new(
352            &self.config,
353            self.evm_opts.clone(),
354            Some(known_contracts),
355            Some(artifact_id.clone()),
356        ));
357        ExecutorBuilder::new()
358            .inspectors(|stack| {
359                stack
360                    .cheatcodes(cheats_config)
361                    .trace_mode(self.trace_mode())
362                    .coverage(self.coverage)
363                    .enable_isolation(self.isolation)
364                    .odyssey(self.odyssey)
365                    .create2_deployer(self.evm_opts.create2_deployer)
366            })
367            .spec_id(self.spec_id)
368            .gas_limit(self.evm_opts.gas_limit())
369            .legacy_assertions(self.config.legacy_assertions)
370            .build(self.env.clone(), db)
371    }
372
373    fn trace_mode(&self) -> TraceMode {
374        TraceMode::default()
375            .with_debug(self.debug)
376            .with_decode_internal(self.decode_internal)
377            .with_verbosity(self.evm_opts.verbosity)
378            .with_state_changes(verbosity() > 4)
379    }
380}
381
382/// Builder used for instantiating the multi-contract runner
383#[derive(Clone, Debug)]
384#[must_use = "builders do nothing unless you call `build` on them"]
385pub struct MultiContractRunnerBuilder {
386    /// The address which will be used to deploy the initial contracts and send all
387    /// transactions
388    pub sender: Option<Address>,
389    /// The initial balance for each one of the deployed smart contracts
390    pub initial_balance: U256,
391    /// The EVM spec to use
392    pub evm_spec: Option<SpecId>,
393    /// The fork to use at launch
394    pub fork: Option<CreateFork>,
395    /// Project config.
396    pub config: Arc<Config>,
397    /// Whether or not to collect coverage info
398    pub coverage: bool,
399    /// Whether or not to collect debug info
400    pub debug: bool,
401    /// Whether to enable steps tracking in the tracer.
402    pub decode_internal: InternalTraceMode,
403    /// Whether to enable call isolation
404    pub isolation: bool,
405    /// Whether to enable Odyssey features.
406    pub odyssey: bool,
407}
408
409impl MultiContractRunnerBuilder {
410    pub fn new(config: Arc<Config>) -> Self {
411        Self {
412            config,
413            sender: Default::default(),
414            initial_balance: Default::default(),
415            evm_spec: Default::default(),
416            fork: Default::default(),
417            coverage: Default::default(),
418            debug: Default::default(),
419            isolation: Default::default(),
420            decode_internal: Default::default(),
421            odyssey: Default::default(),
422        }
423    }
424
425    pub fn sender(mut self, sender: Address) -> Self {
426        self.sender = Some(sender);
427        self
428    }
429
430    pub fn initial_balance(mut self, initial_balance: U256) -> Self {
431        self.initial_balance = initial_balance;
432        self
433    }
434
435    pub fn evm_spec(mut self, spec: SpecId) -> Self {
436        self.evm_spec = Some(spec);
437        self
438    }
439
440    pub fn with_fork(mut self, fork: Option<CreateFork>) -> Self {
441        self.fork = fork;
442        self
443    }
444
445    pub fn set_coverage(mut self, enable: bool) -> Self {
446        self.coverage = enable;
447        self
448    }
449
450    pub fn set_debug(mut self, enable: bool) -> Self {
451        self.debug = enable;
452        self
453    }
454
455    pub fn set_decode_internal(mut self, mode: InternalTraceMode) -> Self {
456        self.decode_internal = mode;
457        self
458    }
459
460    pub fn enable_isolation(mut self, enable: bool) -> Self {
461        self.isolation = enable;
462        self
463    }
464
465    pub fn odyssey(mut self, enable: bool) -> Self {
466        self.odyssey = enable;
467        self
468    }
469
470    /// Given an EVM, proceeds to return a runner which is able to execute all tests
471    /// against that evm
472    pub fn build<C: Compiler<CompilerContract = Contract>>(
473        self,
474        root: &Path,
475        output: &ProjectCompileOutput,
476        env: Env,
477        evm_opts: EvmOpts,
478    ) -> Result<MultiContractRunner> {
479        let contracts = output
480            .artifact_ids()
481            .map(|(id, v)| (id.with_stripped_file_prefixes(root), v))
482            .collect();
483        let linker = Linker::new(root, contracts);
484
485        // Build revert decoder from ABIs of all artifacts.
486        let abis = linker
487            .contracts
488            .iter()
489            .filter_map(|(_, contract)| contract.abi.as_ref().map(|abi| abi.borrow()));
490        let revert_decoder = RevertDecoder::new().with_abis(abis);
491
492        let LinkOutput { libraries, libs_to_deploy } = linker.link_with_nonce_or_address(
493            Default::default(),
494            LIBRARY_DEPLOYER,
495            0,
496            linker.contracts.keys(),
497        )?;
498
499        let linked_contracts = linker.get_linked_artifacts(&libraries)?;
500
501        // Create a mapping of name => (abi, deployment code, Vec<library deployment code>)
502        let mut deployable_contracts = DeployableContracts::default();
503
504        for (id, contract) in linked_contracts.iter() {
505            let Some(abi) = &contract.abi else { continue };
506
507            // if it's a test, link it and add to deployable contracts
508            if abi.constructor.as_ref().map(|c| c.inputs.is_empty()).unwrap_or(true) &&
509                abi.functions().any(|func| func.name.is_any_test())
510            {
511                let Some(bytecode) =
512                    contract.get_bytecode_bytes().map(|b| b.into_owned()).filter(|b| !b.is_empty())
513                else {
514                    continue;
515                };
516
517                deployable_contracts
518                    .insert(id.clone(), TestContract { abi: abi.clone(), bytecode });
519            }
520        }
521
522        let known_contracts = ContractsByArtifact::new(linked_contracts);
523
524        Ok(MultiContractRunner {
525            contracts: deployable_contracts,
526            revert_decoder,
527            known_contracts,
528            libs_to_deploy,
529            libraries,
530
531            fork: self.fork,
532
533            tcfg: TestRunnerConfig {
534                evm_opts,
535                env,
536                spec_id: self.evm_spec.unwrap_or_else(|| self.config.evm_spec_id()),
537                sender: self.sender.unwrap_or(self.config.sender),
538
539                coverage: self.coverage,
540                debug: self.debug,
541                decode_internal: self.decode_internal,
542                inline_config: Arc::new(InlineConfig::new_parsed(output, &self.config)?),
543                isolation: self.isolation,
544                odyssey: self.odyssey,
545
546                config: self.config,
547            },
548        })
549    }
550}
551
552pub fn matches_contract(id: &ArtifactId, abi: &JsonAbi, filter: &dyn TestFilter) -> bool {
553    (filter.matches_path(&id.source) && filter.matches_contract(&id.name)) &&
554        abi.functions().any(|func| is_matching_test(func, filter))
555}
556
557/// Returns `true` if the function is a test function that matches the given filter.
558pub(crate) fn is_matching_test(func: &Function, filter: &dyn TestFilter) -> bool {
559    func.is_any_test() && filter.matches_test(&func.signature())
560}
forge/multi_runner.rs

forge/
multi_runner.rs
