foundry_cheatcodes/evm/
mock.rs

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

foundry_cheatcodes/evm/
mock.rs
