foundry_cheatcodes/test/
assume.rs

Help
1use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Error, Result};
2use alloy_primitives::Address;
3use foundry_evm_core::constants::MAGIC_ASSUME;
4use spec::Vm::{
5    assumeCall, assumeNoRevert_0Call, assumeNoRevert_1Call, assumeNoRevert_2Call, PotentialRevert,
6};
7use std::fmt::Debug;
8
9#[derive(Clone, Debug)]
10pub struct AssumeNoRevert {
11    /// The call depth at which the cheatcode was added.
12    pub depth: u64,
13    /// Acceptable revert parameters for the next call, to be thrown out if they are encountered;
14    /// reverts with parameters not specified here will count as normal reverts and not rejects
15    /// towards the counter.
16    pub reasons: Vec<AcceptableRevertParameters>,
17    /// Address that reverted the call.
18    pub reverted_by: Option<Address>,
19}
20
21/// Parameters for a single anticipated revert, to be thrown out if encountered.
22#[derive(Clone, Debug)]
23pub struct AcceptableRevertParameters {
24    /// The expected revert data returned by the revert
25    pub reason: Vec<u8>,
26    /// If true then only the first 4 bytes of expected data returned by the revert are checked.
27    pub partial_match: bool,
28    /// Contract expected to revert next call.
29    pub reverter: Option<Address>,
30}
31
32impl AcceptableRevertParameters {
33    fn from(potential_revert: &PotentialRevert) -> Self {
34        Self {
35            reason: potential_revert.revertData.to_vec(),
36            partial_match: potential_revert.partialMatch,
37            reverter: if potential_revert.reverter == Address::ZERO {
38                None
39            } else {
40                Some(potential_revert.reverter)
41            },
42        }
43    }
44}
45
46impl Cheatcode for assumeCall {
47    fn apply(&self, _state: &mut Cheatcodes) -> Result {
48        let Self { condition } = self;
49        if *condition {
50            Ok(Default::default())
51        } else {
52            Err(Error::from(MAGIC_ASSUME))
53        }
54    }
55}
56
57impl Cheatcode for assumeNoRevert_0Call {
58    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
59        assume_no_revert(ccx.state, ccx.ecx.journaled_state.depth(), vec![])
60    }
61}
62
63impl Cheatcode for assumeNoRevert_1Call {
64    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
65        let Self { potentialRevert } = self;
66        assume_no_revert(
67            ccx.state,
68            ccx.ecx.journaled_state.depth(),
69            vec![AcceptableRevertParameters::from(potentialRevert)],
70        )
71    }
72}
73
74impl Cheatcode for assumeNoRevert_2Call {
75    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
76        let Self { potentialReverts } = self;
77        assume_no_revert(
78            ccx.state,
79            ccx.ecx.journaled_state.depth(),
80            potentialReverts.iter().map(AcceptableRevertParameters::from).collect(),
81        )
82    }
83}
84
85fn assume_no_revert(
86    state: &mut Cheatcodes,
87    depth: u64,
88    parameters: Vec<AcceptableRevertParameters>,
89) -> Result {
90    ensure!(
91        state.assume_no_revert.is_none(),
92        "you must make another external call prior to calling assumeNoRevert again"
93    );
94
95    state.assume_no_revert = Some(AssumeNoRevert { depth, reasons: parameters, reverted_by: None });
96
97    Ok(Default::default())
98}
foundry_cheatcodes/test/assume.rs

foundry_cheatcodes/test/
assume.rs
