Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: add count to expectCall cheatcodes #4833

Merged
merged 1 commit into from
Apr 27, 2023
Merged

feat: add count to expectCall cheatcodes #4833

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions evm/src/executor/abi/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,9 +102,13 @@ abigen!(
clearMockedCalls()

expectCall(address,bytes)
expectCall(address,bytes,uint64)
expectCall(address,uint256,bytes)
expectCall(address,uint256,bytes,uint64)
expectCall(address,uint256,uint64,bytes)
expectCall(address,uint256,uint64,bytes,uint64)
expectCallMinGas(address,uint256,uint64,bytes)
expectCallMinGas(address,uint256,uint64,bytes,uint64)
expectSafeMemory(uint64,uint64)
expectSafeMemoryCall(uint64,uint64)

Expand Down
124 changes: 99 additions & 25 deletions evm/src/executor/inspector/cheatcodes/expect.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -210,6 +210,8 @@ pub struct ExpectedCallData {
pub gas: Option<u64>,
/// The expected *minimum* gas supplied to the call
pub min_gas: Option<u64>,
/// The number of times the call is expected to be made
pub count: u64,
Comment on lines +213 to +214
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can see how this could be useful

}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
Expand Down Expand Up @@ -293,51 +295,123 @@ pub fn apply<DB: DatabaseExt>(
Ok(Bytes::new())
}
HEVMCalls::ExpectCall0(inner) => {
state.expected_calls.entry(inner.0).or_default().push(ExpectedCallData {
calldata: inner.1.to_vec().into(),
value: None,
gas: None,
min_gas: None,
});
state.expected_calls.entry(inner.0).or_default().push((
ExpectedCallData {
calldata: inner.1.to_vec().into(),
value: None,
gas: None,
min_gas: None,
count: 1,
},
0,
));
Ok(Bytes::new())
}
HEVMCalls::ExpectCall1(inner) => {
state.expected_calls.entry(inner.0).or_default().push(ExpectedCallData {
calldata: inner.2.to_vec().into(),
value: Some(inner.1),
gas: None,
min_gas: None,
});
state.expected_calls.entry(inner.0).or_default().push((
ExpectedCallData {
calldata: inner.1.to_vec().into(),
value: None,
gas: None,
min_gas: None,
count: inner.2,
},
0,
));
Ok(Bytes::new())
}
HEVMCalls::ExpectCall2(inner) => {
state.expected_calls.entry(inner.0).or_default().push((
ExpectedCallData {
calldata: inner.2.to_vec().into(),
value: Some(inner.1),
gas: None,
min_gas: None,
count: 1,
},
0,
));
Ok(Bytes::new())
}
HEVMCalls::ExpectCall3(inner) => {
state.expected_calls.entry(inner.0).or_default().push((
ExpectedCallData {
calldata: inner.2.to_vec().into(),
value: Some(inner.1),
gas: None,
min_gas: None,
count: inner.3,
},
0,
));
Ok(Bytes::new())
}
HEVMCalls::ExpectCall4(inner) => {
let value = inner.1;

// If the value of the transaction is non-zero, the EVM adds a call stipend of 2300 gas
// to ensure that the basic fallback function can be called.
let positive_value_cost_stipend = if value > U256::zero() { 2300 } else { 0 };

state.expected_calls.entry(inner.0).or_default().push(ExpectedCallData {
calldata: inner.3.to_vec().into(),
value: Some(value),
gas: Some(inner.2 + positive_value_cost_stipend),
min_gas: None,
});
state.expected_calls.entry(inner.0).or_default().push((
ExpectedCallData {
calldata: inner.3.to_vec().into(),
value: Some(value),
gas: Some(inner.2 + positive_value_cost_stipend),
min_gas: None,
count: 1,
},
0,
));
Ok(Bytes::new())
}
HEVMCalls::ExpectCall5(inner) => {
let value = inner.1;
let positive_value_cost_stipend = if value > U256::zero() { 2300 } else { 0 };
state.expected_calls.entry(inner.0).or_default().push((
ExpectedCallData {
calldata: inner.3.to_vec().into(),
value: Some(value),
gas: Some(inner.2 + positive_value_cost_stipend),
min_gas: None,
count: inner.4,
},
0,
));
Ok(Bytes::new())
}
HEVMCalls::ExpectCallMinGas(inner) => {
HEVMCalls::ExpectCallMinGas0(inner) => {
let value = inner.1;

// If the value of the transaction is non-zero, the EVM adds a call stipend of 2300 gas
// to ensure that the basic fallback function can be called.
let positive_value_cost_stipend = if value > U256::zero() { 2300 } else { 0 };

state.expected_calls.entry(inner.0).or_default().push(ExpectedCallData {
calldata: inner.3.to_vec().into(),
value: Some(value),
gas: None,
min_gas: Some(inner.2 + positive_value_cost_stipend),
});
state.expected_calls.entry(inner.0).or_default().push((
ExpectedCallData {
calldata: inner.3.to_vec().into(),
value: Some(value),
gas: None,
min_gas: Some(inner.2 + positive_value_cost_stipend),
count: 1,
},
0,
));
Ok(Bytes::new())
}
HEVMCalls::ExpectCallMinGas1(inner) => {
let value = inner.1;
let positive_value_cost_stipend = if value > U256::zero() { 2300 } else { 0 };
state.expected_calls.entry(inner.0).or_default().push((
ExpectedCallData {
calldata: inner.3.to_vec().into(),
value: Some(value),
gas: None,
min_gas: Some(inner.2 + positive_value_cost_stipend),
count: inner.4,
},
0,
));
Ok(Bytes::new())
}
HEVMCalls::MockCall0(inner) => {
Expand Down
53 changes: 28 additions & 25 deletions evm/src/executor/inspector/cheatcodes/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ pub struct Cheatcodes {
pub mocked_calls: BTreeMap<Address, BTreeMap<MockCallDataContext, MockCallReturnData>>,

/// Expected calls
pub expected_calls: BTreeMap<Address, Vec<ExpectedCallData>>,
pub expected_calls: BTreeMap<Address, Vec<(ExpectedCallData, u64)>>,

/// Expected emits
pub expected_emits: Vec<ExpectedEmit>,
Expand Down Expand Up @@ -542,14 +542,14 @@ where
} else if call.contract != HARDHAT_CONSOLE_ADDRESS {
// Handle expected calls
if let Some(expecteds) = self.expected_calls.get_mut(&call.contract) {
if let Some(found_match) = expecteds.iter().position(|expected| {
if let Some((_, count)) = expecteds.iter_mut().find(|(expected, _)| {
expected.calldata.len() <= call.input.len() &&
expected.calldata == call.input[..expected.calldata.len()] &&
expected.value.map_or(true, |value| value == call.transfer.value) &&
expected.gas.map_or(true, |gas| gas == call.gas_limit) &&
expected.min_gas.map_or(true, |min_gas| min_gas <= call.gas_limit)
}) {
expecteds.remove(found_match);
*count += 1;
}
}

Expand Down Expand Up @@ -738,28 +738,31 @@ where

// If the depth is 0, then this is the root call terminating
if data.journaled_state.depth() == 0 {
// Handle expected calls that were not fulfilled
if let Some((address, expecteds)) =
self.expected_calls.iter().find(|(_, expecteds)| !expecteds.is_empty())
{
let ExpectedCallData { calldata, gas, min_gas, value } = &expecteds[0];
let calldata = ethers::types::Bytes::from(calldata.clone());
let expected_values = [
Some(format!("data {calldata}")),
value.map(|v| format!("value {v}")),
gas.map(|g| format!("gas {g}")),
min_gas.map(|g| format!("minimum gas {g}")),
]
.into_iter()
.flatten()
.join(" and ");
return (
Return::Revert,
remaining_gas,
format!("Expected a call to {address:?} with {expected_values}, but got none")
.encode()
.into(),
)
for (address, expecteds) in &self.expected_calls {
for (expected, actual_count) in expecteds {
let ExpectedCallData { calldata, gas, min_gas, value, count } = expected;
let calldata = ethers::types::Bytes::from(calldata.clone());
if *count != *actual_count {
let expected_values = [
Some(format!("data {calldata}")),
value.map(|v| format!("value {v}")),
gas.map(|g| format!("gas {g}")),
min_gas.map(|g| format!("minimum gas {g}")),
]
.into_iter()
.flatten()
.join(" and ");
return (
Return::Revert,
remaining_gas,
format!(
"Expected call to {address:?} with {expected_values} to be called {count} time(s), but was called {actual_count} time(s)"
)
.encode()
.into(),
)
}
}
}

// Check if we have any leftover expected emits
Expand Down
14 changes: 12 additions & 2 deletions forge/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -318,13 +318,23 @@ interface Hevm {
function clearMockedCalls() external;
// Expect a call to an address with the specified calldata.
// Calldata can either be strict or a partial match
function expectCall(address,bytes calldata) external;
function expectCall(address, bytes calldata) external;
// Expect given number of calls to an address with the specified calldata.
// Calldata can either be strict or a partial match
function expectCall(address, bytes calldata, uint64) external;
// Expect a call to an address with the specified msg.value and calldata
function expectCall(address,uint256,bytes calldata) external;
function expectCall(address, uint256, bytes calldata) external;
// Expect a given number of calls to an address with the specified msg.value and calldata
function expectCall(address, uint256, bytes calldata, uint64) external;
// Expect a call to an address with the specified msg.value, gas, and calldata.
function expectCall(address, uint256, uint64, bytes calldata) external;
// Expect a given number of calls to an address with the specified msg.value, gas, and calldata.
function expectCall(address, uint256, uint64, bytes calldata, uint64) external;
// Expect a call to an address with the specified msg.value and calldata, and a *minimum* amount of gas.
function expectCallMinGas(address, uint256, uint64, bytes calldata) external;
// Expect a given number of calls to an address with the specified msg.value and calldata, and a *minimum* amount of gas.
function expectCallMinGas(address, uint256, uint64, bytes calldata, uint64) external;

// Only allows memory writes to offsets [0x00, 0x60) ∪ [min, max) in the current subcontext. If any other
// memory is written to, the test will fail.
function expectSafeMemory(uint64, uint64) external;
Expand Down
13 changes: 13 additions & 0 deletions testdata/cheats/Cheats.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -208,15 +208,28 @@ interface Cheats {
// Calldata can either be strict or a partial match
function expectCall(address, bytes calldata) external;

// Expect given number of calls to an address with the specified calldata.
// Calldata can either be strict or a partial match
function expectCall(address, bytes calldata, uint64) external;

// Expect a call to an address with the specified msg.value and calldata
function expectCall(address, uint256, bytes calldata) external;

// Expect a given number of calls to an address with the specified msg.value and calldata
function expectCall(address, uint256, bytes calldata, uint64) external;

// Expect a call to an address with the specified msg.value, gas, and calldata.
function expectCall(address, uint256, uint64, bytes calldata) external;

// Expect a given number of calls to an address with the specified msg.value, gas, and calldata.
function expectCall(address, uint256, uint64, bytes calldata, uint64) external;

// Expect a call to an address with the specified msg.value and calldata, and a *minimum* amount of gas.
function expectCallMinGas(address, uint256, uint64, bytes calldata) external;

// Expect a given number of calls to an address with the specified msg.value and calldata, and a *minimum* amount of gas.
function expectCallMinGas(address, uint256, uint64, bytes calldata, uint64) external;

// Only allows memory writes to offsets [0x00, 0x60) ∪ [min, max) in the current subcontext. If any other
// memory is written to, the test will fail.
function expectSafeMemory(uint64, uint64) external;
Expand Down
Loading