Enable cheating contract_address in tests internally

Towards #3331

commit-id:1f9d52c8

Author: ddoktorski

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

@@ -155,13 +155,12 @@ pub fn get_cheated_exec_info_ptr(
         }
     }
 
-    if cheated_data.caller_address.is_some() {
-        new_exec_info[2] = MaybeRelocatable::Int(
-            cheated_data
-                .caller_address
-                .expect("No caller address value found for the cheated caller address contract")
-                .into_(),
-        );
+    if let Some(caller_address) = cheated_data.caller_address {
+        new_exec_info[2] = MaybeRelocatable::Int(caller_address.into_());
+    }
+
+    if let Some(contract_address) = cheated_data.contract_address {
+        new_exec_info[3] = MaybeRelocatable::Int(contract_address.into_());
     }
 
     vm.load_data(ptr_cheated_exec_info, &new_exec_info).unwrap();

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/cheatcodes/cheat_execution_info.rs

@@ -59,6 +59,7 @@ pub struct ExecutionInfoMock {
     pub block_info: BlockInfoMock,
     pub tx_info: TxInfoMock,
     pub caller_address: CheatStatus<ContractAddress>,
+    pub contract_address: CheatStatus<ContractAddress>,
 }
 
 #[derive(CairoDeserialize, Clone, Default, Debug)]
@@ -90,11 +91,13 @@ pub struct ExecutionInfoMockOperations {
     pub block_info: BlockInfoMockOperations,
     pub tx_info: TxInfoMockOperations,
     pub caller_address: Operation<ContractAddress>,
+    pub contract_address: Operation<ContractAddress>,
 }
 
 macro_rules! for_all_fields {
     ($macro:ident!) => {
         $macro!(caller_address);
+        $macro!(contract_address);
 
         $macro!(block_info.block_number);
         $macro!(block_info.block_timestamp);

crates/cheatnet/src/state.rs

@@ -353,6 +353,7 @@ pub struct CheatedData {
     pub block_number: Option<u64>,
     pub block_timestamp: Option<u64>,
     pub caller_address: Option<ContractAddress>,
+    pub contract_address: Option<ContractAddress>,
     pub sequencer_address: Option<ContractAddress>,
     pub tx_info: CheatedTxInfo,
 }
@@ -422,6 +423,7 @@ impl CheatnetState {
             block_number: execution_info.block_info.block_number.as_value(),
             block_timestamp: execution_info.block_info.block_timestamp.as_value(),
             caller_address: execution_info.caller_address.as_value(),
+            contract_address: execution_info.contract_address.as_value(),
             sequencer_address: execution_info.block_info.sequencer_address.as_value(),
             tx_info: CheatedTxInfo {
                 version: execution_info.tx_info.version.as_value(),

snforge_std/src/cheatcodes/execution_info.cairo

@@ -4,6 +4,8 @@ use snforge_std::cheatcodes::CheatSpan;
 use crate::cheatcode::execute_cheatcode_and_deserialize;
 
 pub mod caller_address;
+#[doc(hidden)]
+pub mod contract_address;
 pub mod block_number;
 pub mod block_timestamp;
 pub mod sequencer_address;
@@ -123,6 +125,7 @@ struct ExecutionInfoMock {
     block_info: BlockInfoMock,
     tx_info: TxInfoMock,
     caller_address: Operation<ContractAddress>,
+    contract_address: Operation<ContractAddress>,
 }
 
 impl ExecutionInfoMockImpl of Default<ExecutionInfoMock> {
@@ -133,6 +136,7 @@ impl ExecutionInfoMockImpl of Default<ExecutionInfoMock> {
             block_info: Default::default(),
             tx_info: Default::default(),
             caller_address: Operation::Retain,
+            contract_address: Operation::Retain,
         }
     }
 }

snforge_std/src/cheatcodes/execution_info/contract_address.cairo

@@ -0,0 +1,36 @@
+//! These cheatcodes are currently used only internally by the `interact_with_state` cheatcode,
+//! and are specifically intended to cheat on the `TEST_ADDRESS`.
+//! They are not exposed through the public API in a general form, as there are no known use cases
+//! that would require them.
+
+use super::{
+    ExecutionInfoMock, cheat_execution_info, Operation, ContractAddress, CheatArguments, CheatSpan,
+};
+use crate::test_address;
+
+/// Overrides the contract address for the default test address.
+/// After this function is called, any call to `starknet::get_contract_address()`
+/// within the test context will return the provided `contract_address`.
+/// - `contract_address` - The contract address to use for the default test address.
+pub(crate) fn start_cheat_contract_address(contract_address: ContractAddress) {
+    let mut execution_info: ExecutionInfoMock = Default::default();
+
+    execution_info
+        .contract_address =
+            Operation::Start(
+                CheatArguments {
+                    value: contract_address, span: CheatSpan::Indefinite, target: test_address(),
+                },
+            );
+
+    cheat_execution_info(execution_info);
+}
+
+/// Cancels the `start_cheat_contract_address`.
+pub(crate) fn stop_cheat_contract_address() {
+    let mut execution_info: ExecutionInfoMock = Default::default();
+
+    execution_info.contract_address = Operation::Stop(test_address());
+
+    cheat_execution_info(execution_info);
+}