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

foundry_cheatcodes/evm/
mock.rs
