Add expectCall(address,uint256,uint64,bytes) cheatcode

foundry-rs · Feb 23, 2023 · e09e9a8 · e09e9a8
1 parent e049b0d
commit e09e9a8
Show file tree

Hide file tree

Showing 7 changed files with 71 additions and 15 deletions.
diff --git a/evm/src/executor/abi/mod.rs b/evm/src/executor/abi/mod.rs
@@ -81,6 +81,7 @@ ethers::contract::abigen!(
             clearMockedCalls()
             expectCall(address,bytes)
             expectCall(address,uint256,bytes)
+            expectCall(address,uint256,uint64,bytes)
             getCode(string)
             getDeployedCode(string)
             label(address,string)

diff --git a/evm/src/executor/inspector/cheatcodes/expect.rs b/evm/src/executor/inspector/cheatcodes/expect.rs
@@ -206,6 +206,8 @@ pub struct ExpectedCallData {
     pub calldata: Bytes,
     /// The expected value sent in the call
     pub value: Option<U256>,
+    /// The expected gas supplied to the call
+    pub gas: Option<u64>,
 }
 
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
@@ -264,19 +266,33 @@ 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 });
+            state.expected_calls.entry(inner.0).or_default().push(ExpectedCallData {
+                calldata: inner.1.to_vec().into(),
+                value: None,
+                gas: None,
+            });
             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) });
+            state.expected_calls.entry(inner.0).or_default().push(ExpectedCallData {
+                calldata: inner.2.to_vec().into(),
+                value: Some(inner.1),
+                gas: None,
+            });
+            Ok(Bytes::new())
+        }
+        HEVMCalls::ExpectCall2(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),
+            });
             Ok(Bytes::new())
         }
         HEVMCalls::MockCall0(inner) => {

diff --git a/evm/src/executor/inspector/cheatcodes/mod.rs b/evm/src/executor/inspector/cheatcodes/mod.rs
@@ -402,7 +402,8 @@ where
                 if let Some(found_match) = expecteds.iter().position(|expected| {
                     expected.calldata.len() <= call.input.len() &&
                         expected.calldata == call.input[..expected.calldata.len()] &&
-                        expected.value.map(|value| value == call.transfer.value).unwrap_or(true)
+                        expected.value.map(|value| value == call.transfer.value).unwrap_or(true) &&
+                        expected.gas.map(|gas| gas == call.gas_limit).unwrap_or(true)
                 }) {
                     expecteds.remove(found_match);
                 }
@@ -584,10 +585,11 @@ where
                     Return::Revert,
                     remaining_gas,
                     format!(
-                        "Expected a call to {:?} with data {}{}, but got none",
+                        "Expected a call to {:?} with data {}{}{}, but got none",
                         address,
                         ethers::types::Bytes::from(expecteds[0].calldata.clone()),
-                        expecteds[0].value.map(|v| format!(" and value {v}")).unwrap_or_default()
+                        expecteds[0].value.map(|v| format!(" and value {v}")).unwrap_or_default(),
+                        expecteds[0].gas.map(|g| format!(" and gas {g}")).unwrap_or_default(),
                     )
                     .encode()
                     .into(),

diff --git a/forge/README.md b/forge/README.md
@@ -317,6 +317,8 @@ interface Hevm {
     function expectCall(address,bytes calldata) external;
     // Expect a call to an address with the specified msg.value and calldata
     function expectCall(address,uint256,bytes calldata) external;
+    // Expect a call to an address with the specified msg.value, gas, and calldata.
+    function expectCall(address, uint256, uint64, bytes calldata) external;
     // Fetches the contract bytecode from its artifact file
     function getCode(string calldata) external returns (bytes memory);
     // Label an address in test traces

diff --git a/testdata/cache/solidity-files-cache.json b/testdata/cache/solidity-files-cache.json
@@ -509,8 +509,8 @@
       }
     },
     "cheats/ExpectCall.t.sol": {
-      "lastModificationDate": 1661330493212,
-      "contentHash": "c6164b1cfe68bc3bfd2f7171f47fbc75",
+      "lastModificationDate": 1677141892508,
+      "contentHash": "b15f2d663c79b0e205cbf458f81540d6",
       "sourceName": "cheats/ExpectCall.t.sol",
       "solcConfig": {
         "settings": {

diff --git a/testdata/cheats/Cheats.sol b/testdata/cheats/Cheats.sol
@@ -196,6 +196,9 @@ interface Cheats {
     // Expect a call to an address with the specified msg.value and calldata
     function expectCall(address, uint256, bytes calldata) external;
 
+    // Expect a call to an address with the specified msg.value, gas, and calldata.
+    function expectCall(address, uint256, uint64, bytes calldata) external;
+
     // Gets the bytecode from an artifact file. Takes in the relative path to the json file
     function getCode(string calldata) external returns (bytes memory);
 

diff --git a/testdata/cheats/ExpectCall.t.sol b/testdata/cheats/ExpectCall.t.sol
@@ -33,6 +33,14 @@ contract NestedContract {
         return inner.numberA() + inner.numberB();
     }
 
+    function forwardPay() public payable returns (uint256) {
+        return inner.pay{ gas: 50_000, value: 1 }(1);
+    }
+
+    function addHardGasLimit() public view returns (uint256) {
+        return inner.add{ gas: 50_000 }(1, 1);
+    }
+
     function hello() public pure returns (string memory) {
         return "hi";
     }
@@ -104,4 +112,28 @@ contract ExpectCallTest is DSTest {
         cheats.expectCall(address(target), 3, abi.encodeWithSelector(target.pay.selector));
         target.pay{value: 3}(100);
     }
+
+    function testExpectCallWithValueAndGas() public {
+        Contract inner = new Contract();
+        NestedContract target = new NestedContract(inner);
+
+        cheats.expectCall(address(inner), 1, 50_000, abi.encodeWithSelector(inner.pay.selector, 1));
+        target.forwardPay{ value: 1 }();
+    }
+
+    function testExpectCallWithNoValueAndGas() public {
+        Contract inner = new Contract();
+        NestedContract target = new NestedContract(inner);
+
+        cheats.expectCall(address(inner), 0, 50_000, abi.encodeWithSelector(inner.add.selector, 1, 1));
+        target.addHardGasLimit();
+    }
+
+    function testFailExpectCallWithNoValueAndWrongGas() public {
+        Contract inner = new Contract();
+        NestedContract target = new NestedContract(inner);
+
+        cheats.expectCall(address(inner), 0, 25_000, abi.encodeWithSelector(inner.add.selector, 1, 1));
+        target.addHardGasLimit();
+    }
 }