snforge_std: revert breaking changes on mock_call

Add `mock_call_when`, `start_mock_call_when` and `stop_mock_call_when`

Author: ptisserand

crates/forge/tests/integration/cheat_fork.rs

@@ -155,7 +155,7 @@ fn mock_call_cairo0_contract() {
     let test = test_case!(formatdoc!(
         r#"
             use starknet::{{contract_address_const}};
-            use snforge_std::{{start_mock_call, stop_mock_call, MockCallData}};
+            use snforge_std::{{start_mock_call, stop_mock_call}};
 
             #[starknet::interface]
             trait IERC20<TContractState> {{
@@ -173,11 +173,11 @@ fn mock_call_cairo0_contract() {
 
                 assert(eth_dispatcher.name() == 'Ether', 'invalid name');
 
-                start_mock_call(eth_dispatcher.contract_address, selector!("name"), MockCallData::Any, 'NotEther');
+                start_mock_call(eth_dispatcher.contract_address, selector!("name"), 'NotEther');
 
                 assert(eth_dispatcher.name() == 'NotEther', 'invalid mocked name');
 
-                stop_mock_call(eth_dispatcher.contract_address, selector!("name"), MockCallData::Any);
+                stop_mock_call(eth_dispatcher.contract_address, selector!("name"));
 
                 assert(eth_dispatcher.name() == 'Ether', 'invalid name after mock');
             }}

crates/forge/tests/integration/mock_call.rs

@@ -10,7 +10,7 @@ fn mock_call_simple() {
         indoc!(
             r#"
         use result::ResultTrait;
-        use snforge_std::{ declare, ContractClassTrait, DeclareResultTrait, start_mock_call, stop_mock_call, MockCallData };
+        use snforge_std::{ declare, ContractClassTrait, DeclareResultTrait, start_mock_call, stop_mock_call };
 
         #[starknet::interface]
         trait IMockChecker<TContractState> {
@@ -26,19 +26,13 @@ fn mock_call_simple() {
 
             let dispatcher = IMockCheckerDispatcher { contract_address };
 
-            let specific_mock_ret_data = 421;
-            let default_mock_ret_data = 404;
-            let expected_calldata = MockCallData::Values([].span());
-            start_mock_call(contract_address, selector!("get_thing"), expected_calldata, specific_mock_ret_data);
-            start_mock_call(contract_address, selector!("get_thing"), MockCallData::Any, default_mock_ret_data);
-            let thing = dispatcher.get_thing();
-            assert(thing == specific_mock_ret_data, 'Incorrect thing');
+            let mock_ret_data = 421;
 
-            stop_mock_call(contract_address, selector!("get_thing"), expected_calldata);
+            start_mock_call(contract_address, selector!("get_thing"), mock_ret_data);
             let thing = dispatcher.get_thing();
-            assert(thing == default_mock_ret_data, 'Incorrect thing');
+            assert(thing == 421, 'Incorrect thing');
 
-            stop_mock_call(contract_address, selector!("get_thing"), MockCallData::Any);
+            stop_mock_call(contract_address, selector!("get_thing"));
             let thing = dispatcher.get_thing();
             assert(thing == 420, 'Incorrect thing');
         }
@@ -51,12 +45,12 @@ fn mock_call_simple() {
             let (contract_address, _) = contract.deploy(@calldata).unwrap();
 
             let mock_ret_data = 421;
-            start_mock_call(contract_address, selector!("get_thing"), MockCallData::Any, mock_ret_data);
+            start_mock_call(contract_address, selector!("get_thing"), mock_ret_data);
 
             let dispatcher = IMockCheckerDispatcher { contract_address };
             let thing = dispatcher.get_thing();
 
-            assert(thing == 421, 'Incorrect thing all catch');
+            assert(thing == 421, 'Incorrect thing');
         }
     "#
         ),
@@ -79,7 +73,7 @@ fn mock_call_complex_types() {
         use result::ResultTrait;
         use array::ArrayTrait;
         use serde::Serde;
-        use snforge_std::{ declare, ContractClassTrait, DeclareResultTrait, start_mock_call, MockCallData };
+        use snforge_std::{ declare, ContractClassTrait, DeclareResultTrait, start_mock_call };
 
         #[starknet::interface]
         trait IMockChecker<TContractState> {
@@ -103,7 +97,7 @@ fn mock_call_complex_types() {
             let dispatcher = IMockCheckerDispatcher { contract_address };
 
             let mock_ret_data = StructThing {item_one: 412, item_two: 421};
-            start_mock_call(contract_address, selector!("get_struct_thing"), MockCallData::Any, mock_ret_data);
+            start_mock_call(contract_address, selector!("get_struct_thing"), mock_ret_data);
 
             let thing: StructThing = dispatcher.get_struct_thing();
 
@@ -121,7 +115,7 @@ fn mock_call_complex_types() {
             let dispatcher = IMockCheckerDispatcher { contract_address };
 
             let mock_ret_data =  array![ StructThing {item_one: 112, item_two: 121}, StructThing {item_one: 412, item_two: 421} ];
-            start_mock_call(contract_address, selector!("get_arr_thing"), MockCallData::Any, mock_ret_data);
+            start_mock_call(contract_address, selector!("get_arr_thing"), mock_ret_data);
 
             let things: Array<StructThing> = dispatcher.get_arr_thing();
 
@@ -152,7 +146,7 @@ fn mock_calls() {
         indoc!(
             r#"
         use result::ResultTrait;
-        use snforge_std::{ declare, ContractClassTrait, DeclareResultTrait, mock_call, MockCallData };
+        use snforge_std::{ declare, ContractClassTrait, DeclareResultTrait, mock_call, start_mock_call, stop_mock_call };
 
         #[starknet::interface]
         trait IMockChecker<TContractState> {
@@ -170,7 +164,7 @@ fn mock_calls() {
 
             let mock_ret_data = 421;
 
-            mock_call(contract_address, selector!("get_thing"), MockCallData::Any, mock_ret_data, 1);
+            mock_call(contract_address, selector!("get_thing"), mock_ret_data, 1);
 
             let thing = dispatcher.get_thing();
             assert_eq!(thing, 421);
@@ -190,7 +184,7 @@ fn mock_calls() {
 
             let mock_ret_data = 421;
 
-            mock_call(contract_address, selector!("get_thing"), MockCallData::Any, mock_ret_data, 2);
+            mock_call(contract_address, selector!("get_thing"), mock_ret_data, 2);
 
             let thing = dispatcher.get_thing();
             assert_eq!(thing, 421);
@@ -213,3 +207,184 @@ fn mock_calls() {
     let result = run_test_case(&test);
     assert_passed(&result);
 }
+
+#[test]
+fn mock_call_when_simple() {
+    let test = test_case!(
+        indoc!(
+            r#"
+        use result::ResultTrait;
+        use snforge_std::{ declare, ContractClassTrait, DeclareResultTrait, start_mock_call_when, stop_mock_call_when, MockCallData };
+
+        #[starknet::interface]
+        trait IMockChecker<TContractState> {
+            fn get_thing(ref self: TContractState) -> felt252;
+        }
+
+        #[test]
+        fn mock_call_when_simple() {
+            let calldata = array![420];
+
+            let contract = declare("MockChecker").unwrap().contract_class();
+            let (contract_address, _) = contract.deploy(@calldata).unwrap();
+
+            let dispatcher = IMockCheckerDispatcher { contract_address };
+            
+            let specific_mock_ret_data = 421;
+            let default_mock_ret_data = 404;
+            let expected_calldata = MockCallData::Values([].span());
+
+            start_mock_call_when(contract_address, selector!("get_thing"), expected_calldata, specific_mock_ret_data);
+            start_mock_call_when(contract_address, selector!("get_thing"), MockCallData::Any, default_mock_ret_data);
+            let thing = dispatcher.get_thing();
+            assert(thing == specific_mock_ret_data, 'Incorrect thing');
+
+            stop_mock_call_when(contract_address, selector!("get_thing"), expected_calldata);
+            let thing = dispatcher.get_thing();
+            assert(thing == default_mock_ret_data, 'Incorrect thing');
+
+            stop_mock_call_when(contract_address, selector!("get_thing"), MockCallData::Any);
+            let thing = dispatcher.get_thing();
+            assert(thing == 420, 'Incorrect thing');
+        }
+
+        #[test]
+        fn mock_call_when_simple_before_dispatcher_created() {
+            let calldata = array![420];
+
+            let contract = declare("MockChecker").unwrap().contract_class();
+            let (contract_address, _) = contract.deploy(@calldata).unwrap();
+
+            let specific_mock_ret_data = 421;
+            let default_mock_ret_data = 404;
+            let expected_calldata = MockCallData::Values([].span());
+            
+            start_mock_call_when(contract_address, selector!("get_thing"), expected_calldata, specific_mock_ret_data);
+            start_mock_call_when(contract_address, selector!("get_thing"), MockCallData::Any, default_mock_ret_data);
+            let dispatcher = IMockCheckerDispatcher { contract_address };
+            let thing = dispatcher.get_thing();
+
+            assert(thing == specific_mock_ret_data, 'Incorrect thing');
+
+            stop_mock_call_when(contract_address, selector!("get_thing"), expected_calldata);
+            let thing = dispatcher.get_thing();
+            assert(thing == default_mock_ret_data, 'Incorrect thing');
+
+            stop_mock_call_when(contract_address, selector!("get_thing"), MockCallData::Any);
+            let thing = dispatcher.get_thing();
+            assert(thing == 420, 'Incorrect thing');
+        }
+    "#
+        ),
+        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_call_when_complex_types() {
+    let test = test_case!(
+        indoc!(
+            r#"
+        use result::ResultTrait;
+        use array::ArrayTrait;
+        use serde::Serde;
+        use snforge_std::{ declare, ContractClassTrait, DeclareResultTrait, start_mock_call_when, stop_mock_call_when, MockCallData };
+
+        #[starknet::interface]
+        trait IMockChecker<TContractState> {
+            fn get_struct_thing(ref self: TContractState) -> StructThing;
+            fn get_arr_thing(ref self: TContractState) -> Array<StructThing>;
+        }
+
+        #[derive(Serde, Drop)]
+        struct StructThing {
+            item_one: felt252,
+            item_two: felt252,
+        }
+
+        #[test]
+        fn start_mock_call_when_return_struct() {
+            let calldata = array![420];
+
+            let contract = declare("MockChecker").unwrap().contract_class();
+            let (contract_address, _) = contract.deploy(@calldata).unwrap();
+
+            let dispatcher = IMockCheckerDispatcher { contract_address };
+
+            let default_mock_ret_data = StructThing {item_one: 412, item_two: 421};
+            let specific_mock_ret_data = StructThing {item_one: 404, item_two: 401};
+            let expected_calldata = MockCallData::Values([].span());
+
+            start_mock_call_when(contract_address, selector!("get_struct_thing"), MockCallData::Any, default_mock_ret_data);
+            start_mock_call_when(contract_address, selector!("get_struct_thing"), expected_calldata, specific_mock_ret_data);
+
+            let thing: StructThing = dispatcher.get_struct_thing();
+
+            assert(thing.item_one == 404, 'thing.item_one');
+            assert(thing.item_two == 401, 'thing.item_two');
+
+            stop_mock_call_when(contract_address, selector!("get_struct_thing"), expected_calldata);
+            let thing: StructThing = dispatcher.get_struct_thing();
+
+            assert(thing.item_one == 412, 'thing.item_one');
+            assert(thing.item_two == 421, 'thing.item_two');
+        }
+
+        #[test]
+        fn start_mock_call_when_return_arr() {
+            let calldata = array![420];
+
+            let contract = declare("MockChecker").unwrap().contract_class();
+            let (contract_address, _) = contract.deploy(@calldata).unwrap();
+
+            let dispatcher = IMockCheckerDispatcher { contract_address };
+
+            let default_mock_ret_data =  array![ StructThing {item_one: 112, item_two: 121}, StructThing {item_one: 412, item_two: 421} ];
+            let specific_mock_ret_data =  array![ StructThing {item_one: 212, item_two: 221}, StructThing {item_one: 512, item_two: 521} ];
+
+            let expected_calldata = MockCallData::Values([].span());
+
+            start_mock_call_when(contract_address, selector!("get_arr_thing"), MockCallData::Any, default_mock_ret_data);
+            start_mock_call_when(contract_address, selector!("get_arr_thing"), expected_calldata, specific_mock_ret_data);
+
+            let things: Array<StructThing> = dispatcher.get_arr_thing();
+
+            let thing = things.at(0);
+            assert(*thing.item_one == 212, 'thing1.item_one 1');
+            assert(*thing.item_two == 221, 'thing1.item_two');
+
+            let thing = things.at(1);
+            assert(*thing.item_one == 512, 'thing2.item_one 2');
+            assert(*thing.item_two == 521, 'thing2.item_two');
+
+            stop_mock_call_when(contract_address, selector!("get_arr_thing"), expected_calldata);
+
+            let things: Array<StructThing> = dispatcher.get_arr_thing();
+
+            let thing = things.at(0);
+            assert(*thing.item_one == 112, 'thing1.item_one 3');
+            assert(*thing.item_two == 121, 'thing1.item_two');
+
+            let thing = things.at(1);
+            assert(*thing.item_one == 412, 'thing2.item_one 4');
+            assert(*thing.item_two == 421, 'thing2.item_two');
+        }
+    "#
+        ),
+        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);
+}

crates/forge/tests/integration/test_state.rs

@@ -678,7 +678,7 @@ fn inconsistent_syscall_pointers() {
         r#"
         use starknet::ContractAddress;
         use starknet::info::get_block_number;
-        use snforge_std::{start_mock_call, MockCallData};
+        use snforge_std::start_mock_call;
 
         #[starknet::interface]
         trait IContract<TContractState> {
@@ -689,7 +689,7 @@ fn inconsistent_syscall_pointers() {
         fn inconsistent_syscall_pointers() {
             // verifies if SyscallHandler.syscal_ptr is incremented correctly when calling a contract
             let address = 'address'.try_into().unwrap();
-            start_mock_call(address, selector!("get_value"), MockCallData::Any, 55);
+            start_mock_call(address, selector!("get_value"), 55);
             let contract = IContractDispatcher { contract_address: address };
             contract.get_value(address);
             get_block_number();