starknet_os_flow_tests: localize expected storage updates in test_os_logic

Author: dorimedini-starkware

crates/starknet_os_flow_tests/src/tests.rs

@@ -12,9 +12,9 @@ use starknet_api::core::{
     calculate_contract_address,
     ClassHash,
     CompiledClassHash as StarknetAPICompiledClassHash,
+    ContractAddress,
     EthAddress,
     Nonce,
-    PatriciaKey,
 };
 use starknet_api::executable_transaction::{
     DeclareTransaction,
@@ -48,7 +48,7 @@ use starknet_api::transaction::{
     MessageToL1,
     TransactionVersion,
 };
-use starknet_api::{calldata, declare_tx_args, invoke_tx_args, storage_key};
+use starknet_api::{calldata, declare_tx_args, invoke_tx_args};
 use starknet_committer::block_committer::input::{
     StarknetStorageKey,
     StarknetStorageValue,
@@ -377,6 +377,20 @@ async fn test_os_logic(#[values(1, 3)] n_blocks_in_multi_block: usize) {
     let (mut test_manager, mut nonce_manager, _) =
         TestManager::<DictStateReader>::new_with_default_initial_state([]).await;
     let n_expected_txs = 29;
+    let mut expected_storage_updates: HashMap<
+        ContractAddress,
+        HashMap<StarknetStorageKey, StarknetStorageValue>,
+    > = HashMap::new();
+    let mut update_expected_storage = |address: ContractAddress, key: Felt, value: Felt| {
+        let key = StarknetStorageKey(StorageKey(key.try_into().unwrap()));
+        let value = StarknetStorageValue(value);
+        expected_storage_updates
+            .entry(address)
+            .and_modify(|map| {
+                map.insert(key, value);
+            })
+            .or_insert_with(|| HashMap::from([(key, value)]));
+    };
 
     // Declare a Cairo 0 test contract.
     let cairo0_test_contract = FeatureContract::TestContract(CairoVersion::Cairo0);
@@ -411,6 +425,14 @@ async fn test_os_logic(#[values(1, 3)] n_blocks_in_multi_block: usize) {
         );
         contract_addresses.push(address);
         test_manager.add_invoke_tx(deploy_tx, None);
+        // Update expected storage diff, if the ctor calldata writes a nonzero value.
+        if ctor_calldata[1] != 0 {
+            update_expected_storage(
+                address,
+                Felt::from(ctor_calldata[0]),
+                Felt::from(ctor_calldata[1]),
+            );
+        }
     }
 
     // Call set_value(address=85, value=47) on the first contract.
@@ -424,6 +446,7 @@ async fn test_os_logic(#[values(1, 3)] n_blocks_in_multi_block: usize) {
         resource_bounds: *NON_TRIVIAL_RESOURCE_BOUNDS,
     };
     test_manager.add_invoke_tx_from_args(invoke_tx_args, &CHAIN_ID_FOR_TESTS, None);
+    update_expected_storage(contract_addresses[0], Felt::from(85), Felt::from(47));
 
     // Call set_value(address=81, value=0) on the first contract.
     // Used to test redundant value update (0 -> 0) and make sure it is not written to on-chain
@@ -458,6 +481,7 @@ async fn test_os_logic(#[values(1, 3)] n_blocks_in_multi_block: usize) {
         resource_bounds: *NON_TRIVIAL_RESOURCE_BOUNDS,
     };
     test_manager.add_invoke_tx_from_args(invoke_tx_args, &CHAIN_ID_FOR_TESTS, None);
+    update_expected_storage(contract_addresses[1], Felt::from(15), Felt::ONE);
 
     let invoke_tx_args = invoke_tx_args! {
         sender_address: *FUNDED_ACCOUNT_ADDRESS,
@@ -612,6 +636,11 @@ async fn test_os_logic(#[values(1, 3)] n_blocks_in_multi_block: usize) {
         resource_bounds: *NON_TRIVIAL_RESOURCE_BOUNDS,
     };
     test_manager.add_invoke_tx_from_args(invoke_tx_args, &CHAIN_ID_FOR_TESTS, None);
+    update_expected_storage(
+        delegate_proxy_address,
+        **get_storage_var_address("implementation_hash", &[]),
+        test_class_hash.0,
+    );
 
     // Call test_get_contract_address(expected_address=delegate_proxy_address) through the delegate
     // proxy.
@@ -635,6 +664,7 @@ async fn test_os_logic(#[values(1, 3)] n_blocks_in_multi_block: usize) {
         resource_bounds: *NON_TRIVIAL_RESOURCE_BOUNDS,
     };
     test_manager.add_invoke_tx_from_args(invoke_tx_args, &CHAIN_ID_FOR_TESTS, None);
+    update_expected_storage(delegate_proxy_address, Felt::from(123), Felt::from(456));
 
     // Call test_get_caller_address(expected_address=account_address) through the delegate proxy.
     let invoke_tx_args = invoke_tx_args! {
@@ -690,6 +720,14 @@ async fn test_os_logic(#[values(1, 3)] n_blocks_in_multi_block: usize) {
         nonce: l1_handler_nonce,
         selector: l1_handler_selector,
     };
+    update_expected_storage(
+        delegate_proxy_address,
+        **get_storage_var_address(
+            "two_counters",
+            &[Felt::from(expected_message_to_l2.from_address)],
+        ),
+        *expected_message_to_l2.payload.0.last().unwrap(),
+    );
 
     // Call test_library_call_syntactic_sugar from contract_addresses[0] to test library calls
     // using the syntactic sugar of 'library_call_<FUNCTION>'.
@@ -704,18 +742,31 @@ async fn test_os_logic(#[values(1, 3)] n_blocks_in_multi_block: usize) {
         resource_bounds: *NON_TRIVIAL_RESOURCE_BOUNDS,
     };
     test_manager.add_invoke_tx_from_args(invoke_tx_args, &CHAIN_ID_FOR_TESTS, None);
+    update_expected_storage(contract_addresses[0], Felt::from(444), Felt::from(666));
 
     // Call add_signature_to_counters(index=2021).
+    let index = Felt::from(2021);
+    let signature = TransactionSignature(Arc::new(vec![Felt::from(100), Felt::from(200)]));
     let invoke_tx_args = invoke_tx_args! {
         sender_address: *FUNDED_ACCOUNT_ADDRESS,
         nonce: nonce_manager.next(*FUNDED_ACCOUNT_ADDRESS),
         calldata: create_calldata(
-            contract_addresses[0], "add_signature_to_counters", &[Felt::from(2021)]
+            contract_addresses[0], "add_signature_to_counters", &[index]
         ),
         resource_bounds: *NON_TRIVIAL_RESOURCE_BOUNDS,
-        signature: TransactionSignature(Arc::new(vec![Felt::from(100), Felt::from(200)])),
+        signature: signature.clone(),
     };
     test_manager.add_invoke_tx_from_args(invoke_tx_args, &CHAIN_ID_FOR_TESTS, None);
+    update_expected_storage(
+        contract_addresses[0],
+        **get_storage_var_address("two_counters", &[index]),
+        signature.0[0],
+    );
+    update_expected_storage(
+        contract_addresses[0],
+        **get_storage_var_address("two_counters", &[index]) + Felt::ONE,
+        signature.0[1],
+    );
 
     // Declare test_contract2.
     let test_contract2 = FeatureContract::TestContract2;
@@ -751,6 +802,7 @@ async fn test_os_logic(#[values(1, 3)] n_blocks_in_multi_block: usize) {
         signature: TransactionSignature(Arc::new(vec![Felt::from(100)])),
     };
     test_manager.add_invoke_tx_from_args(invoke_tx_args, &CHAIN_ID_FOR_TESTS, None);
+    update_expected_storage(contract_addresses[1], Felt::from(555), Felt::from(888));
 
     // Use library_call_l1_handler to invoke test_contract2.test_l1_handler_storage_write with
     // from_address=85, address=666, value=999.
@@ -773,6 +825,7 @@ async fn test_os_logic(#[values(1, 3)] n_blocks_in_multi_block: usize) {
         signature: TransactionSignature(Arc::new(vec![Felt::from(100)])),
     };
     test_manager.add_invoke_tx_from_args(invoke_tx_args, &CHAIN_ID_FOR_TESTS, None);
+    update_expected_storage(contract_addresses[1], Felt::from(666), Felt::from(999));
 
     // Replace the class of contract_addresses[0] to the class of test_contract2.
     let invoke_tx_args = invoke_tx_args! {
@@ -785,63 +838,6 @@ async fn test_os_logic(#[values(1, 3)] n_blocks_in_multi_block: usize) {
     };
     test_manager.add_invoke_tx_from_args(invoke_tx_args, &CHAIN_ID_FOR_TESTS, None);
 
-    // Expected values:
-
-    // Storage updates:
-    let storage_updates = HashMap::from([
-        (
-            contract_addresses[0],
-            HashMap::from([
-                (StarknetStorageKey(storage_key!(85u16)), StarknetStorageValue(Felt::from(47))),
-                (StarknetStorageKey(storage_key!(321u16)), StarknetStorageValue(Felt::from(543))),
-                (StarknetStorageKey(storage_key!(444u16)), StarknetStorageValue(Felt::from(666))),
-                (
-                    StarknetStorageKey(get_storage_var_address(
-                        "two_counters",
-                        &[Felt::from(2021)],
-                    )),
-                    StarknetStorageValue(Felt::from(100)),
-                ),
-                (
-                    StarknetStorageKey(StorageKey(
-                        PatriciaKey::try_from(
-                            **get_storage_var_address("two_counters", &[Felt::from(2021)])
-                                + Felt::ONE,
-                        )
-                        .unwrap(),
-                    )),
-                    StarknetStorageValue(Felt::from(200)),
-                ),
-            ]),
-        ),
-        (
-            contract_addresses[1],
-            HashMap::from([
-                (StarknetStorageKey(storage_key!(15u16)), StarknetStorageValue(Felt::from(1))),
-                (StarknetStorageKey(storage_key!(111u16)), StarknetStorageValue(Felt::from(987))),
-                (StarknetStorageKey(storage_key!(555u16)), StarknetStorageValue(Felt::from(888))),
-                (StarknetStorageKey(storage_key!(666u16)), StarknetStorageValue(Felt::from(999))),
-            ]),
-        ),
-        (
-            delegate_proxy_address,
-            HashMap::from([
-                (StarknetStorageKey(storage_key!(123u16)), StarknetStorageValue(Felt::from(456))),
-                (
-                    StarknetStorageKey(get_storage_var_address("implementation_hash", &[])),
-                    StarknetStorageValue(test_class_hash.0),
-                ),
-                (
-                    StarknetStorageKey(get_storage_var_address(
-                        "two_counters",
-                        &[Felt::from(expected_message_to_l2.from_address)],
-                    )),
-                    StarknetStorageValue(*expected_message_to_l2.payload.0.last().unwrap()),
-                ),
-            ]),
-        ),
-    ]);
-
     // Expected number of txs.
     assert_eq!(test_manager.total_txs(), n_expected_txs);
 
@@ -857,6 +853,7 @@ async fn test_os_logic(#[values(1, 3)] n_blocks_in_multi_block: usize) {
 
     // Perform validations.
     let perform_global_validations = true;
-    let partial_state_diff = Some(&StateDiff { storage_updates, ..Default::default() });
+    let partial_state_diff =
+        Some(&StateDiff { storage_updates: expected_storage_updates, ..Default::default() });
     test_output.perform_validations(perform_global_validations, partial_state_diff);
 }