chore(execution): fix the tests using create_state_changes_for_test f…

…or readability and correctness

crates/blockifier/src/state/cached_state_test.rs

@@ -3,6 +3,7 @@ use std::collections::HashMap;
 use assert_matches::assert_matches;
 use indexmap::indexmap;
 use pretty_assertions::assert_eq;
+use rstest::rstest;
 use starknet_api::core::PatriciaKey;
 use starknet_api::hash::StarkHash;
 use starknet_api::{class_hash, contract_address, patricia_key, stark_felt};
@@ -13,6 +14,8 @@ use crate::test_utils::cached_state::deprecated_create_test_state;
 use crate::test_utils::dict_state_reader::DictStateReader;
 use crate::test_utils::{get_test_contract_class, TEST_CLASS_HASH, TEST_EMPTY_CONTRACT_CLASS_HASH};
 
+pub const CONTRACT_ADDRESS: &str = "0x100";
+
 fn set_initial_state_values(
     state: &mut CachedState<DictStateReader>,
     class_hash_to_class: ContractClassMapping,
@@ -257,9 +260,10 @@ fn cached_state_state_diff_conversion() {
 
 fn create_state_changes_for_test<S: StateReader>(
     state: &mut CachedState<S>,
+    sender_address: Option<ContractAddress>,
     fee_token_address: ContractAddress,
 ) -> StateChanges {
-    let contract_address = contract_address!("0x100");
+    let contract_address = contract_address!(CONTRACT_ADDRESS);
     let contract_address2 = contract_address!("0x101");
     let class_hash = class_hash!("0x10");
     let compiled_class_hash = CompiledClassHash(stark_felt!("0x11"));
@@ -277,36 +281,46 @@ fn create_state_changes_for_test<S: StateReader>(
     // As the second access:
     state.set_storage_at(contract_address, key, storage_val).unwrap();
 
-    // Return the resulting state changes.
-    state
-        .get_actual_state_changes_for_fee_charge(fee_token_address, Some(contract_address))
-        .unwrap()
+    if let Some(sender_address) = sender_address {
+        // Charge fee from the sender.
+        let sender_balance_key = get_fee_token_var_address(sender_address);
+        state.set_storage_at(fee_token_address, sender_balance_key, stark_felt!("0x1999")).unwrap();
+    }
+
+    state.get_actual_state_changes_for_fee_charge(fee_token_address, sender_address).unwrap()
 }
 
-#[test]
-fn test_get_actual_state_changes_for_fee_charge() {
+#[rstest]
+fn test_get_actual_state_changes_for_fee_charge(
+    #[values(Some(contract_address!("0x102")), None)] sender_address: Option<ContractAddress>,
+) {
     let mut state: CachedState<DictStateReader> = CachedState::default();
-    let state_changes = create_state_changes_for_test(&mut state, contract_address!("0x17"));
-    assert_eq!(
-        StateChangesCount::from(&state_changes),
-        StateChangesCount {
-            n_storage_updates: 2, // 1 for storage update + 1 for sender balance update.
-            n_modified_contracts: 2,
-            n_class_hash_updates: 1,
-            n_compiled_class_hash_updates: 1
-        }
-    );
+    let fee_token_address = contract_address!("0x17");
+    let state_changes =
+        create_state_changes_for_test(&mut state, sender_address, fee_token_address);
+
+    let expected_state_changes_count = StateChangesCount {
+        // 1 for storage update + 1 for sender balance update if sender is defined.
+        n_storage_updates: 1 + usize::from(sender_address.is_some()),
+        n_class_hash_updates: 1,
+        n_compiled_class_hash_updates: 1,
+        n_modified_contracts: 2,
+    };
+    assert_eq!(StateChangesCount::from(&state_changes), expected_state_changes_count);
 }
 
-#[test]
-fn test_state_changes_merge() {
+#[rstest]
+fn test_state_changes_merge(
+    #[values(Some(contract_address!("0x102")), None)] sender_address: Option<ContractAddress>,
+) {
     // Create a transactional state containing the `create_state_changes_for_test` logic, get the
     // state changes and then commit.
     let mut state: CachedState<DictStateReader> = CachedState::default();
     let mut transactional_state = CachedState::create_transactional(&mut state);
     let block_context = BlockContext::create_for_testing();
     let fee_token_address = block_context.chain_info.fee_token_addresses.eth_fee_token_address;
-    let state_changes1 = create_state_changes_for_test(&mut transactional_state, fee_token_address);
+    let state_changes1 =
+        create_state_changes_for_test(&mut transactional_state, sender_address, fee_token_address);
     transactional_state.commit();
 
     // After performing `commit`, the transactional state is moved (into state).  We need to create
@@ -330,18 +344,32 @@ fn test_state_changes_merge() {
 
     // Get the storage updates addresses and keys from the state_changes1, to overwrite.
     let mut storage_updates_keys = state_changes1.storage_updates.keys();
-    let (contract_address, storage_key) = *storage_updates_keys.next().unwrap();
-    let (contract_address2, storage_key2) = *storage_updates_keys.next().unwrap();
+    // THis is unstable!!!
+    let (contract_address, storage_key) = if storage_updates_keys.len() == 1 {
+        *storage_updates_keys.next().unwrap()
+    } else if storage_updates_keys.len() == 2 {
+        let (address1, storage_key1) = *storage_updates_keys.next().unwrap();
+        let (address2, storage_key2) = *storage_updates_keys.next().unwrap();
+        if address1 == contract_address!(CONTRACT_ADDRESS) {
+            (address1, storage_key1)
+        } else if address2 == contract_address!(CONTRACT_ADDRESS) {
+            (address2, storage_key2)
+        } else {
+            panic!(
+                "Unexpected storage updates addresses. The storage updates are: {:?}",
+                state_changes1.storage_updates
+            );
+        }
+    } else {
+        panic!("Unexpected number of storage updates. {:?}", state_changes1.storage_updates);
+    };
     // A new address, not included in state_changes1, to write to.
     let new_contract_address = ContractAddress(patricia_key!("0x111"));
 
     // Overwrite existing and new storage values.
     transactional_state
         .set_storage_at(contract_address, storage_key, stark_felt!("0x1234"))
         .unwrap();
-    transactional_state
-        .set_storage_at(contract_address2, storage_key2, stark_felt!("0x4321"))
-        .unwrap();
     transactional_state
         .set_storage_at(new_contract_address, storage_key, stark_felt!("0x43210"))
         .unwrap();
@@ -356,7 +384,7 @@ fn test_state_changes_merge() {
     // states. We expect the state_changes to match the merged state_changes of the transactional
     // states, but only when done in the right order.
     let state_changes_final =
-        state.get_actual_state_changes_for_fee_charge(fee_token_address, None).unwrap();
+        state.get_actual_state_changes_for_fee_charge(fee_token_address, sender_address).unwrap();
     assert_eq!(
         StateChanges::merge(vec![
             state_changes1.clone(),