Add tests for cheat_account_contract_address cheatcode (#3722)

<!-- Reference any GitHub issues resolved by this PR -->

Advances #2145

## Introduced changes

<!-- A brief description of the changes -->

- Adds tests for `cheat_account_contract_address` cheatcode

## Checklist

<!-- Make sure all of these are complete -->

- [x] Linked relevant issue
- [ ] Updated relevant documentation
- [x] Added relevant tests
- [x] Performed self-review of the code
- [ ] Added changes to `CHANGELOG.md`

---------

Co-authored-by: ddoktorski <[email protected]>

Author: naijauser

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

@@ -0,0 +1,76 @@
+use crate::{
+    runtime_extensions::forge_runtime_extension::cheatcodes::cheat_execution_info::{
+        CheatArguments, ExecutionInfoMockOperations, Operation, TxInfoMockOperations,
+    },
+    state::{CheatSpan, CheatnetState},
+};
+use starknet_api::core::ContractAddress;
+
+impl CheatnetState {
+    pub fn cheat_account_contract_address(
+        &mut self,
+        target: ContractAddress,
+        account_contract_address: ContractAddress,
+        span: CheatSpan,
+    ) {
+        self.cheat_execution_info(ExecutionInfoMockOperations {
+            tx_info: TxInfoMockOperations {
+                account_contract_address: Operation::Start(CheatArguments {
+                    value: account_contract_address.into(),
+                    span,
+                    target,
+                }),
+                ..Default::default()
+            },
+            ..Default::default()
+        });
+    }
+
+    pub fn start_cheat_account_contract_address(
+        &mut self,
+        target: ContractAddress,
+        account_contract_address: ContractAddress,
+    ) {
+        self.cheat_account_contract_address(
+            target,
+            account_contract_address,
+            CheatSpan::Indefinite,
+        );
+    }
+
+    pub fn stop_cheat_account_contract_address(&mut self, target: ContractAddress) {
+        self.cheat_execution_info(ExecutionInfoMockOperations {
+            tx_info: TxInfoMockOperations {
+                account_contract_address: Operation::Stop(target),
+                ..Default::default()
+            },
+            ..Default::default()
+        });
+    }
+
+    pub fn start_cheat_account_contract_address_global(
+        &mut self,
+        account_contract_address: ContractAddress,
+    ) {
+        self.cheat_execution_info(ExecutionInfoMockOperations {
+            tx_info: TxInfoMockOperations {
+                account_contract_address: Operation::StartGlobal(account_contract_address.into()),
+                ..Default::default()
+            },
+            ..Default::default()
+        });
+    }
+
+    pub fn stop_cheat_account_contract_address_global(
+        &mut self,
+        account_contract_address: ContractAddress,
+    ) {
+        self.cheat_execution_info(ExecutionInfoMockOperations {
+            tx_info: TxInfoMockOperations {
+                account_contract_address: Operation::StartGlobal(account_contract_address.into()),
+                ..Default::default()
+            },
+            ..Default::default()
+        });
+    }
+}

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

@@ -3,6 +3,7 @@ use cairo_vm::vm::errors::hint_errors::HintError;
 use runtime::EnhancedHintError;
 use starknet_types_core::felt::Felt;
 
+pub mod cheat_account_contract_address;
 pub mod cheat_block_hash;
 pub mod cheat_block_number;
 pub mod cheat_block_timestamp;

crates/cheatnet/tests/cheatcodes/cheat_account_contract_address.rs

@@ -0,0 +1,104 @@
+use std::num::NonZeroUsize;
+
+use crate::{cheatcodes::test_environment::TestEnvironment, common::assertions::assert_success};
+use cheatnet::state::CheatSpan;
+use starknet_api::core::ContractAddress;
+use starknet_types_core::felt::Felt;
+
+trait CheatAccountContractAddressTrait {
+    fn cheat_account_contract_address(
+        &mut self,
+        target: ContractAddress,
+        new_address: u128,
+        span: CheatSpan,
+    );
+    fn start_cheat_account_contract_address(&mut self, target: ContractAddress, new_address: u128);
+    fn stop_cheat_account_contract_address(&mut self, target: ContractAddress);
+}
+
+impl CheatAccountContractAddressTrait for TestEnvironment {
+    fn cheat_account_contract_address(
+        &mut self,
+        target: ContractAddress,
+        new_address: u128,
+        span: CheatSpan,
+    ) {
+        self.cheatnet_state.cheat_account_contract_address(
+            target,
+            ContractAddress::from(new_address),
+            span,
+        );
+    }
+
+    fn start_cheat_account_contract_address(&mut self, target: ContractAddress, new_address: u128) {
+        self.cheatnet_state
+            .start_cheat_account_contract_address(target, ContractAddress::from(new_address));
+    }
+
+    fn stop_cheat_account_contract_address(&mut self, target: ContractAddress) {
+        self.cheatnet_state
+            .stop_cheat_account_contract_address(target);
+    }
+}
+
+#[test]
+fn cheat_account_contract_address_simple() {
+    let mut test_env = TestEnvironment::new();
+
+    let contract_address = test_env.deploy("CheatTxInfoChecker", &[]);
+
+    let output = test_env.call_contract(&contract_address, "get_account_contract_address", &[]);
+    assert_success(output, &[Felt::from(0)]);
+
+    test_env.start_cheat_account_contract_address(contract_address, 123);
+
+    let output = test_env.call_contract(&contract_address, "get_account_contract_address", &[]);
+    assert_success(output, &[Felt::from(123)]);
+}
+
+#[test]
+fn cheat_account_contract_address_stop() {
+    let mut test_env = TestEnvironment::new();
+
+    let contract_address = test_env.deploy("CheatTxInfoChecker", &[]);
+
+    test_env.start_cheat_account_contract_address(contract_address, 123);
+
+    assert_success(
+        test_env.call_contract(&contract_address, "get_account_contract_address", &[]),
+        &[Felt::from(123)],
+    );
+
+    test_env.stop_cheat_account_contract_address(contract_address);
+
+    assert_success(
+        test_env.call_contract(&contract_address, "get_account_contract_address", &[]),
+        &[Felt::from(0)],
+    );
+}
+
+#[test]
+fn cheat_account_contract_address_simple_with_span() {
+    let mut test_env = TestEnvironment::new();
+
+    let contract_address = test_env.deploy("CheatTxInfoChecker", &[]);
+
+    test_env.cheat_account_contract_address(
+        contract_address,
+        123,
+        CheatSpan::TargetCalls(NonZeroUsize::new(2).unwrap()),
+    );
+
+    assert_success(
+        test_env.call_contract(&contract_address, "get_account_contract_address", &[]),
+        &[Felt::from(123)],
+    );
+    assert_success(
+        test_env.call_contract(&contract_address, "get_account_contract_address", &[]),
+        &[Felt::from(123)],
+    );
+    assert_success(
+        test_env.call_contract(&contract_address, "get_account_contract_address", &[]),
+        &[Felt::from(0)],
+    );
+}

crates/cheatnet/tests/cheatcodes/mod.rs

@@ -1,5 +1,6 @@
 mod test_environment;
 
+mod cheat_account_contract_address;
 mod cheat_block_hash;
 mod cheat_block_number;
 mod cheat_block_timestamp;

crates/cheatnet/tests/contracts/src/cheat_tx_info/tx_info_checker.cairo

@@ -1,14 +1,17 @@
+use starknet::ContractAddress;
+
 #[starknet::interface]
 trait ICheatTxInfoChecker<TContractState> {
+    fn get_account_contract_address(ref self: TContractState) -> ContractAddress;
     fn get_transaction_hash(self: @TContractState) -> felt252;
     fn get_tx_hash_and_emit_event(ref self: TContractState) -> felt252;
     fn get_tx_info(self: @TContractState) -> starknet::TxInfo;
 }
 
 #[starknet::contract]
 mod CheatTxInfoChecker {
-    use starknet::SyscallResultTrait;
     use starknet::syscalls::get_execution_info_v2_syscall;
+    use starknet::{ContractAddress, SyscallResultTrait, get_tx_info};
 
     #[storage]
     struct Storage {}
@@ -40,5 +43,10 @@ mod CheatTxInfoChecker {
             let execution_info = get_execution_info_v2_syscall().unwrap_syscall().unbox();
             execution_info.tx_info.unbox()
         }
+
+        fn get_account_contract_address(ref self: ContractState) -> ContractAddress {
+            let tx_info = get_tx_info();
+            tx_info.account_contract_address
+        }
     }
 }