fix(cheatnet): default to any calldata entry when specific CheatSpan::TargetCalls is 0

Author: ptisserand

crates/cheatnet/src/runtime_extensions/call_to_blockifier_runtime_extension/execution/entry_point.rs

@@ -2,7 +2,7 @@ use std::cell::RefCell;
 use super::cairo1_execution::execute_entry_point_call_cairo1;
 use crate::runtime_extensions::call_to_blockifier_runtime_extension::execution::deprecated::cairo0_execution::execute_entry_point_call_cairo0;
 use crate::runtime_extensions::call_to_blockifier_runtime_extension::CheatnetState;
-use crate::state::{CallTrace, CallTraceNode, CheatStatus, EncounteredError};
+use crate::state::{CallTrace, CallTraceNode, CheatSpan, CheatStatus, EncounteredError};
 use blockifier::execution::call_info::{CallExecution, Retdata};
 use blockifier::{
     execution::{
@@ -278,11 +278,14 @@ fn get_mocked_function_cheat_status<'a>(
         Some(contract_functions) => {
             let call_data_hash = poseidon_hash_many(call.calldata.0.iter());
             let key = (call.entry_point_selector, call_data_hash);
-            if contract_functions.contains_key(&key) {
-                contract_functions.get_mut(&key)
-            } else {
-                let key_zero = (call.entry_point_selector, Felt::zero());
-                contract_functions.get_mut(&key_zero)
+            let key_zero = (call.entry_point_selector, Felt::zero());
+
+            match contract_functions.get(&key) {
+                Some(CheatStatus::Cheated(_, CheatSpan::TargetCalls(0))) => {
+                    contract_functions.get_mut(&key_zero)
+                }
+                Some(CheatStatus::Cheated(_, _)) => contract_functions.get_mut(&key),
+                _ => contract_functions.get_mut(&key_zero),
             }
         }
     }

crates/forge/tests/integration/mock_call.rs

@@ -388,3 +388,195 @@ fn mock_call_when_complex_types() {
     let result = run_test_case(&test);
     assert_passed(&result);
 }
+
+#[test]
+fn mock_calls_when() {
+    let test = test_case!(
+        indoc!(
+            r#"
+        use result::ResultTrait;
+        use snforge_std::{ declare, ContractClassTrait, DeclareResultTrait, mock_call_when, MockCallData};
+
+        #[starknet::interface]
+        trait IMockChecker<TContractState> {
+            fn get_thing(ref self: TContractState) -> felt252;
+        }
+
+        #[test]
+        fn mock_call_when_one_specific() {
+            let calldata = array![420];
+
+            let contract = declare("MockChecker").unwrap().contract_class();
+            let (contract_address, _) = contract.deploy(@calldata).unwrap();
+
+            let dispatcher = IMockCheckerDispatcher { contract_address };
+
+            let mock_ret_data = 421;
+            let expected_calldata = MockCallData::Values([].span());
+            mock_call_when(contract_address, selector!("get_thing"), expected_calldata, mock_ret_data, 1);
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 421);
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 420);
+        }
+
+        #[test]
+        fn mock_call_when_twice_specific() {
+            let calldata = array![420];
+
+            let contract = declare("MockChecker").unwrap().contract_class();
+            let (contract_address, _) = contract.deploy(@calldata).unwrap();
+
+            let dispatcher = IMockCheckerDispatcher { contract_address };
+
+            let mock_ret_data = 421;
+            let expected_calldata = MockCallData::Values([].span());
+            mock_call_when(contract_address, selector!("get_thing"), expected_calldata, mock_ret_data, 2);
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 421);
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 421);
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 420);
+        }
+
+        #[test]
+        fn mock_call_when_one_any() {
+            let calldata = array![420];
+
+            let contract = declare("MockChecker").unwrap().contract_class();
+            let (contract_address, _) = contract.deploy(@calldata).unwrap();
+
+            let dispatcher = IMockCheckerDispatcher { contract_address };
+
+            let mock_ret_data = 421;
+            mock_call_when(contract_address, selector!("get_thing"), MockCallData::Any, mock_ret_data, 1);
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 421);
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 420);
+        }
+
+        #[test]
+        fn mock_call_when_twice_any() {
+            let calldata = array![420];
+
+            let contract = declare("MockChecker").unwrap().contract_class();
+            let (contract_address, _) = contract.deploy(@calldata).unwrap();
+
+            let dispatcher = IMockCheckerDispatcher { contract_address };
+
+            let mock_ret_data = 421;
+            mock_call_when(contract_address, selector!("get_thing"), MockCallData::Any, mock_ret_data, 2);
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 421);
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 421);
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 420);
+        }
+
+    "#
+        ),
+        Contract::from_code_path(
+            "MockChecker".to_string(),
+            Path::new("tests/data/contracts/mock_checker.cairo"),
+        )
+        .unwrap()
+    );
+
+    let result = run_test_case(&test);
+    assert_passed(&result);
+}
+
+#[test]
+fn mock_calls_when_mixed() {
+    let test = test_case!(
+        indoc!(
+            r#"
+        use result::ResultTrait;
+        use snforge_std::{ declare, ContractClassTrait, DeclareResultTrait, mock_call_when, MockCallData};
+
+        #[starknet::interface]
+        trait IMockChecker<TContractState> {
+            fn get_thing(ref self: TContractState) -> felt252;
+        }
+
+        #[test]
+        fn mock_call_when_one() {
+            let calldata = array![420];
+
+            let contract = declare("MockChecker").unwrap().contract_class();
+            let (contract_address, _) = contract.deploy(@calldata).unwrap();
+
+            let dispatcher = IMockCheckerDispatcher { contract_address };
+
+            let mock_ret_data = 421;
+            let expected_calldata = MockCallData::Values([].span());
+            mock_call_when(contract_address, selector!("get_thing"), expected_calldata, mock_ret_data, 1);
+            mock_call_when(contract_address, selector!("get_thing"), MockCallData::Any, 422, 1);
+            
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 421, "Specific calldata");
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 422, "Any calldata");
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 420);
+        }
+
+        #[test]
+        fn mock_call_when_multi() {
+            let calldata = array![420];
+
+            let contract = declare("MockChecker").unwrap().contract_class();
+            let (contract_address, _) = contract.deploy(@calldata).unwrap();
+
+            let dispatcher = IMockCheckerDispatcher { contract_address };
+
+            let mock_ret_data = 421;
+            let expected_calldata = MockCallData::Values([].span());
+            mock_call_when(contract_address, selector!("get_thing"), expected_calldata, mock_ret_data, 3);
+            mock_call_when(contract_address, selector!("get_thing"), MockCallData::Any, 422, 2);
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 421, "1st Specific calldata");
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 421, "2nd Specific calldata");
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 421, "3rd Specific calldata");
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 422, "1st Any calldata");
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 422, "2nd Any calldata");
+
+            let thing = dispatcher.get_thing();
+            assert_eq!(thing, 420);
+        }
+    "#
+        ),
+        Contract::from_code_path(
+            "MockChecker".to_string(),
+            Path::new("tests/data/contracts/mock_checker.cairo"),
+        )
+        .unwrap()
+    );
+
+    let result = run_test_case(&test);
+    assert_passed(&result);
+}