foundry_cheatcodes/evm/
mock.rs

1use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Result, Vm::*};
2use alloy_primitives::{Address, Bytes, U256};
3use revm::{bytecode::Bytecode, context::JournalTr, interpreter::InstructionResult};
4use std::{cmp::Ordering, collections::VecDeque};
5
6/// Mocked call data.
7#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
8pub struct MockCallDataContext {
9    /// The partial calldata to match for mock
10    pub calldata: Bytes,
11    /// The value to match for mock
12    pub value: Option<U256>,
13}
14
15/// Mocked return data.
16#[derive(Clone, Debug)]
17pub struct MockCallReturnData {
18    /// The return type for the mocked call
19    pub ret_type: InstructionResult,
20    /// Return data or error
21    pub data: Bytes,
22}
23
24impl PartialOrd for MockCallDataContext {
25    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
26        Some(self.cmp(other))
27    }
28}
29
30impl Ord for MockCallDataContext {
31    fn cmp(&self, other: &Self) -> Ordering {
32        // Calldata matching is reversed to ensure that a tighter match is
33        // returned if an exact match is not found. In case, there is
34        // a partial match to calldata that is more specific than
35        // a match to a msg.value, then the more specific calldata takes
36        // precedence.
37        self.calldata.cmp(&other.calldata).reverse().then(self.value.cmp(&other.value).reverse())
38    }
39}
40
41impl Cheatcode for clearMockedCallsCall {
42    fn apply(&self, state: &mut Cheatcodes) -> Result {
43        let Self {} = self;
44        state.mocked_calls = Default::default();
45        Ok(Default::default())
46    }
47}
48
49impl Cheatcode for mockCall_0Call {
50    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
51        let Self { callee, data, returnData } = self;
52        let _ = make_acc_non_empty(callee, ccx)?;
53
54        mock_call(ccx.state, callee, data, None, returnData, InstructionResult::Return);
55        Ok(Default::default())
56    }
57}
58
59impl Cheatcode for mockCall_1Call {
60    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
61        let Self { callee, msgValue, data, returnData } = self;
62        let _ = make_acc_non_empty(callee, ccx)?;
63
64        mock_call(ccx.state, callee, data, Some(msgValue), returnData, InstructionResult::Return);
65        Ok(Default::default())
66    }
67}
68
69impl Cheatcode for mockCall_2Call {
70    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
71        let Self { callee, data, returnData } = self;
72        let _ = make_acc_non_empty(callee, ccx)?;
73
74        mock_call(
75            ccx.state,
76            callee,
77            &Bytes::from(*data),
78            None,
79            returnData,
80            InstructionResult::Return,
81        );
82        Ok(Default::default())
83    }
84}
85
86impl Cheatcode for mockCall_3Call {
87    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
88        let Self { callee, msgValue, data, returnData } = self;
89        let _ = make_acc_non_empty(callee, ccx)?;
90
91        mock_call(
92            ccx.state,
93            callee,
94            &Bytes::from(*data),
95            Some(msgValue),
96            returnData,
97            InstructionResult::Return,
98        );
99        Ok(Default::default())
100    }
101}
102
103impl Cheatcode for mockCalls_0Call {
104    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
105        let Self { callee, data, returnData } = self;
106        let _ = make_acc_non_empty(callee, ccx)?;
107
108        mock_calls(ccx.state, callee, data, None, returnData, InstructionResult::Return);
109        Ok(Default::default())
110    }
111}
112
113impl Cheatcode for mockCalls_1Call {
114    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
115        let Self { callee, msgValue, data, returnData } = self;
116        let _ = make_acc_non_empty(callee, ccx)?;
117
118        mock_calls(ccx.state, callee, data, Some(msgValue), returnData, InstructionResult::Return);
119        Ok(Default::default())
120    }
121}
122
123impl Cheatcode for mockCallRevert_0Call {
124    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
125        let Self { callee, data, revertData } = self;
126        let _ = make_acc_non_empty(callee, ccx)?;
127
128        mock_call(ccx.state, callee, data, None, revertData, InstructionResult::Revert);
129        Ok(Default::default())
130    }
131}
132
133impl Cheatcode for mockCallRevert_1Call {
134    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
135        let Self { callee, msgValue, data, revertData } = self;
136        let _ = make_acc_non_empty(callee, ccx)?;
137
138        mock_call(ccx.state, callee, data, Some(msgValue), revertData, InstructionResult::Revert);
139        Ok(Default::default())
140    }
141}
142
143impl Cheatcode for mockCallRevert_2Call {
144    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
145        let Self { callee, data, revertData } = self;
146        let _ = make_acc_non_empty(callee, ccx)?;
147
148        mock_call(
149            ccx.state,
150            callee,
151            &Bytes::from(*data),
152            None,
153            revertData,
154            InstructionResult::Revert,
155        );
156        Ok(Default::default())
157    }
158}
159
160impl Cheatcode for mockCallRevert_3Call {
161    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
162        let Self { callee, msgValue, data, revertData } = self;
163        let _ = make_acc_non_empty(callee, ccx)?;
164
165        mock_call(
166            ccx.state,
167            callee,
168            &Bytes::from(*data),
169            Some(msgValue),
170            revertData,
171            InstructionResult::Revert,
172        );
173        Ok(Default::default())
174    }
175}
176
177impl Cheatcode for mockFunctionCall {
178    fn apply(&self, state: &mut Cheatcodes) -> Result {
179        let Self { callee, target, data } = self;
180        state.mocked_functions.entry(*callee).or_default().insert(data.clone(), *target);
181
182        Ok(Default::default())
183    }
184}
185
186fn mock_call(
187    state: &mut Cheatcodes,
188    callee: &Address,
189    cdata: &Bytes,
190    value: Option<&U256>,
191    rdata: &Bytes,
192    ret_type: InstructionResult,
193) {
194    mock_calls(state, callee, cdata, value, std::slice::from_ref(rdata), ret_type)
195}
196
197fn mock_calls(
198    state: &mut Cheatcodes,
199    callee: &Address,
200    cdata: &Bytes,
201    value: Option<&U256>,
202    rdata_vec: &[Bytes],
203    ret_type: InstructionResult,
204) {
205    state.mocked_calls.entry(*callee).or_default().insert(
206        MockCallDataContext { calldata: cdata.clone(), value: value.copied() },
207        rdata_vec
208            .iter()
209            .map(|rdata| MockCallReturnData { ret_type, data: rdata.clone() })
210            .collect::<VecDeque<_>>(),
211    );
212}
213
214// Etches a single byte onto the account if it is empty to circumvent the `extcodesize`
215// check Solidity might perform.
216fn make_acc_non_empty(callee: &Address, ecx: &mut CheatsCtxt) -> Result {
217    let acc = ecx.journaled_state.load_account(*callee)?;
218
219    let empty_bytecode = acc.info.code.as_ref().is_none_or(Bytecode::is_empty);
220    if empty_bytecode {
221        let code = Bytecode::new_raw(Bytes::from_static(&[0u8]));
222        ecx.journaled_state.set_code(*callee, code);
223    }
224
225    Ok(Default::default())
226}
foundry_cheatcodes/evm/mock.rs

foundry_cheatcodes/evm/
mock.rs
