foundry_cheatcodes/evm/
mock.rs

1use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Result, Vm::*};
2use alloy_primitives::{Address, Bytes, U256};
3use foundry_evm_core::evm::FoundryEvmNetwork;
4use revm::{
5    bytecode::Bytecode,
6    context::{ContextTr, JournalTr},
7    interpreter::InstructionResult,
8};
9use std::{
10    cmp::Ordering,
11    collections::{BTreeMap, VecDeque},
12};
13
14/// Mocked call data.
15#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
16pub struct MockCallDataContext {
17    /// The partial calldata to match for mock
18    pub calldata: Bytes,
19    /// The value to match for mock
20    pub value: Option<U256>,
21}
22
23/// Mocked return data.
24#[derive(Clone, Debug)]
25pub struct MockCallReturnData {
26    /// The return type for the mocked call
27    pub ret_type: InstructionResult,
28    /// Return data or error
29    pub data: Bytes,
30}
31
32impl PartialOrd for MockCallDataContext {
33    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
34        Some(self.cmp(other))
35    }
36}
37
38impl Ord for MockCallDataContext {
39    fn cmp(&self, other: &Self) -> Ordering {
40        // Calldata matching is reversed to ensure that a tighter match is
41        // returned if an exact match is not found. In case, there is
42        // a partial match to calldata that is more specific than
43        // a match to a msg.value, then the more specific calldata takes
44        // precedence.
45        self.calldata.cmp(&other.calldata).reverse().then(self.value.cmp(&other.value).reverse())
46    }
47}
48
49impl Cheatcode for clearMockedCallsCall {
50    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
51        let Self {} = self;
52        state.mocked_calls = Default::default();
53        Ok(Default::default())
54    }
55}
56
57impl Cheatcode for mockCall_0Call {
58    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
59        let Self { callee, data, returnData } = self;
60        let _ = make_acc_non_empty(callee, ccx)?;
61
62        mock_call(ccx.state, callee, data, None, returnData, InstructionResult::Return);
63        Ok(Default::default())
64    }
65}
66
67impl Cheatcode for mockCall_1Call {
68    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
69        let Self { callee, msgValue, data, returnData } = self;
70        let _ = make_acc_non_empty(callee, ccx)?;
71
72        mock_call(ccx.state, callee, data, Some(msgValue), returnData, InstructionResult::Return);
73        Ok(Default::default())
74    }
75}
76
77impl Cheatcode for mockCall_2Call {
78    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
79        let Self { callee, data, returnData } = self;
80        let _ = make_acc_non_empty(callee, ccx)?;
81
82        mock_call(
83            ccx.state,
84            callee,
85            &Bytes::from(*data),
86            None,
87            returnData,
88            InstructionResult::Return,
89        );
90        Ok(Default::default())
91    }
92}
93
94impl Cheatcode for mockCall_3Call {
95    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
96        let Self { callee, msgValue, data, returnData } = self;
97        let _ = make_acc_non_empty(callee, ccx)?;
98
99        mock_call(
100            ccx.state,
101            callee,
102            &Bytes::from(*data),
103            Some(msgValue),
104            returnData,
105            InstructionResult::Return,
106        );
107        Ok(Default::default())
108    }
109}
110
111impl Cheatcode for mockCall_4Call {
112    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
113        let Self { callee, data, returnData, injectCode } = self;
114        if *injectCode {
115            let _ = make_acc_non_empty(callee, ccx)?;
116        }
117
118        mock_call(ccx.state, callee, data, None, returnData, InstructionResult::Return);
119        Ok(Default::default())
120    }
121}
122
123impl Cheatcode for mockCalls_0Call {
124    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
125        let Self { callee, data, returnData } = self;
126        let _ = make_acc_non_empty(callee, ccx)?;
127
128        mock_calls(ccx.state, callee, data, None, returnData, InstructionResult::Return);
129        Ok(Default::default())
130    }
131}
132
133impl Cheatcode for mockCalls_1Call {
134    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
135        let Self { callee, msgValue, data, returnData } = self;
136        let _ = make_acc_non_empty(callee, ccx)?;
137
138        mock_calls(ccx.state, callee, data, Some(msgValue), returnData, InstructionResult::Return);
139        Ok(Default::default())
140    }
141}
142
143impl Cheatcode for mockCallRevert_0Call {
144    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
145        let Self { callee, data, revertData } = self;
146        let _ = make_acc_non_empty(callee, ccx)?;
147
148        mock_call(ccx.state, callee, data, None, revertData, InstructionResult::Revert);
149        Ok(Default::default())
150    }
151}
152
153impl Cheatcode for mockCallRevert_1Call {
154    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
155        let Self { callee, msgValue, data, revertData } = self;
156        let _ = make_acc_non_empty(callee, ccx)?;
157
158        mock_call(ccx.state, callee, data, Some(msgValue), revertData, InstructionResult::Revert);
159        Ok(Default::default())
160    }
161}
162
163impl Cheatcode for mockCallRevert_2Call {
164    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
165        let Self { callee, data, revertData } = self;
166        let _ = make_acc_non_empty(callee, ccx)?;
167
168        mock_call(
169            ccx.state,
170            callee,
171            &Bytes::from(*data),
172            None,
173            revertData,
174            InstructionResult::Revert,
175        );
176        Ok(Default::default())
177    }
178}
179
180impl Cheatcode for mockCallRevert_3Call {
181    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
182        let Self { callee, msgValue, data, revertData } = self;
183        let _ = make_acc_non_empty(callee, ccx)?;
184
185        mock_call(
186            ccx.state,
187            callee,
188            &Bytes::from(*data),
189            Some(msgValue),
190            revertData,
191            InstructionResult::Revert,
192        );
193        Ok(Default::default())
194    }
195}
196
197impl Cheatcode for mockFunctionCall {
198    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
199        let Self { callee, target, data } = self;
200        state.mocked_functions.entry(*callee).or_default().insert(data.clone(), *target);
201
202        Ok(Default::default())
203    }
204}
205
206fn mock_call<FEN: FoundryEvmNetwork>(
207    state: &mut Cheatcodes<FEN>,
208    callee: &Address,
209    cdata: &Bytes,
210    value: Option<&U256>,
211    rdata: &Bytes,
212    ret_type: InstructionResult,
213) {
214    mock_calls(state, callee, cdata, value, std::slice::from_ref(rdata), ret_type)
215}
216
217fn mock_calls<FEN: FoundryEvmNetwork>(
218    state: &mut Cheatcodes<FEN>,
219    callee: &Address,
220    cdata: &Bytes,
221    value: Option<&U256>,
222    rdata_vec: &[Bytes],
223    ret_type: InstructionResult,
224) {
225    state.mocked_calls.entry(*callee).or_default().insert(
226        MockCallDataContext { calldata: cdata.clone(), value: value.copied() },
227        rdata_vec
228            .iter()
229            .map(|rdata| MockCallReturnData { ret_type, data: rdata.clone() })
230            .collect::<VecDeque<_>>(),
231    );
232}
233
234// Etches a single byte onto the account if it is empty to circumvent the `extcodesize`
235// check Solidity might perform.
236fn make_acc_non_empty<FEN: FoundryEvmNetwork>(
237    callee: &Address,
238    ccx: &mut CheatsCtxt<'_, '_, FEN>,
239) -> Result {
240    let empty_bytecode = {
241        let acc = ccx.ecx.journal_mut().load_account(*callee)?;
242        acc.info.code.as_ref().is_none_or(Bytecode::is_empty)
243    };
244    if empty_bytecode {
245        let code = Bytecode::new_raw(Bytes::from_static(&[0u8]));
246        ccx.ecx.journal_mut().set_code(*callee, code);
247    }
248
249    Ok(Default::default())
250}
251
252/// Finds the return data queue of the mock matching a call.
253///
254/// An exact calldata and value match wins. Otherwise, the first mock in map order whose calldata
255/// prefixes `input` and whose value, if set, equals `value` is used.
256pub(crate) fn find_mock_returns<'a, T>(
257    mocks: &'a mut BTreeMap<MockCallDataContext, VecDeque<T>>,
258    input: &Bytes,
259    value: Option<U256>,
260) -> Option<&'a mut VecDeque<T>> {
261    let ctx = MockCallDataContext { calldata: input.clone(), value };
262    // Reversed `Ord` puts all matches at or after `ctx`, with the exact key first.
263    mocks
264        .range_mut(ctx..)
265        .find(|(mock, _)| {
266            input.get(..mock.calldata.len()) == Some(&mock.calldata[..])
267                && mock.value.is_none_or(|mock_value| Some(mock_value) == value)
268        })
269        .map(|(_, v)| v)
270}
271
272/// Consumes the front return data of a mock, keeping the last one for every later call.
273pub(crate) fn advance_mock_returns<T>(queue: &mut VecDeque<T>) {
274    if queue.len() > 1 {
275        queue.pop_front();
276    }
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282
283    fn mock(
284        calldata: &'static [u8],
285        value: Option<u64>,
286        result: u8,
287    ) -> (MockCallDataContext, VecDeque<u8>) {
288        (
289            MockCallDataContext {
290                calldata: Bytes::from_static(calldata),
291                value: value.map(U256::from),
292            },
293            VecDeque::from([result]),
294        )
295    }
296
297    #[test]
298    fn mock_matching_precedence() {
299        let mut mocks = BTreeMap::from([
300            mock(b"call", Some(1), 1),
301            mock(b"call", Some(4), 9),
302            mock(b"call", None, 2),
303            mock(b"cal", Some(1), 3),
304            mock(b"cal", Some(0), 4),
305            mock(b"cal", Some(2), 5),
306            mock(b"cal", None, 6),
307            mock(b"ca", Some(3), 7),
308            mock(b"call-longer", Some(1), 8),
309        ]);
310
311        for (input, value, expected) in [
312            (b"call", Some(1), 1),
313            (b"call", Some(3), 2),
314            (b"call", None, 2),
315            (b"calx", Some(1), 3),
316            (b"calx", Some(3), 6),
317        ] {
318            let queue =
319                find_mock_returns(&mut mocks, &Bytes::from_static(input), value.map(U256::from))
320                    .unwrap();
321            assert_eq!(queue.front(), Some(&expected), "input: {input:?}, value: {value:?}");
322        }
323    }
324
325    #[test]
326    fn absent_transfer_value_does_not_match_zero() {
327        let mut mocks = BTreeMap::from([mock(b"call", Some(0), 1)]);
328        let input = Bytes::from_static(b"call");
329        assert!(find_mock_returns(&mut mocks, &input, None).is_none());
330        assert_eq!(
331            find_mock_returns(&mut mocks, &input, Some(U256::ZERO)).unwrap().front(),
332            Some(&1)
333        );
334
335        mocks.extend([mock(b"cal", None, 2)]);
336        assert_eq!(find_mock_returns(&mut mocks, &input, None).unwrap().front(), Some(&2));
337    }
338
339    #[test]
340    fn empty_matched_queue_is_distinct_from_no_match() {
341        let mut mocks = BTreeMap::from([mock(b"cal", None, 1)]);
342        mocks.insert(
343            MockCallDataContext { calldata: Bytes::from_static(b"call"), value: None },
344            VecDeque::new(),
345        );
346
347        let queue = find_mock_returns(&mut mocks, &Bytes::from_static(b"call"), None).unwrap();
348        assert!(queue.is_empty());
349        assert!(find_mock_returns(&mut mocks, &Bytes::from_static(b"other"), None).is_none());
350    }
351
352    #[test]
353    fn advancing_returns_keeps_last_result() {
354        let mut queue = VecDeque::from([1, 2, 3]);
355        advance_mock_returns(&mut queue);
356        assert_eq!(queue, VecDeque::from([2, 3]));
357        advance_mock_returns(&mut queue);
358        assert_eq!(queue, VecDeque::from([3]));
359        advance_mock_returns(&mut queue);
360        assert_eq!(queue, VecDeque::from([3]));
361
362        queue.clear();
363        advance_mock_returns(&mut queue);
364        assert!(queue.is_empty());
365    }
366}
foundry_cheatcodes/evm/mock.rs

foundry_cheatcodes/evm/
mock.rs
