From 7ede3d202f8d41a1f7798e9a69124a6733dfaa41 Mon Sep 17 00:00:00 2001
From: meship-starkware <meshi@starkware.co>
Date: Thu, 30 May 2024 09:46:29 +0300
Subject: [PATCH] test(concurrency): add utilitis for commit tx test

---
 .../src/concurrency/worker_logic_test.rs      | 46 +++++++++++++++++--
 crates/blockifier/src/fee/fee_utils.rs        | 14 ++++--
 2 files changed, 53 insertions(+), 7 deletions(-)

diff --git a/crates/blockifier/src/concurrency/worker_logic_test.rs b/crates/blockifier/src/concurrency/worker_logic_test.rs
index 197e682f32..d25d6e1182 100644
--- a/crates/blockifier/src/concurrency/worker_logic_test.rs
+++ b/crates/blockifier/src/concurrency/worker_logic_test.rs
@@ -2,7 +2,7 @@ use std::collections::HashMap;
 use std::sync::Mutex;
 
 use num_bigint::BigUint;
-use starknet_api::core::{ContractAddress, PatriciaKey};
+use starknet_api::core::{ContractAddress, Nonce, PatriciaKey};
 use starknet_api::hash::{StarkFelt, StarkHash};
 use starknet_api::transaction::{ContractAddressSalt, Fee, TransactionVersion};
 use starknet_api::{contract_address, patricia_key, stark_felt};
@@ -23,9 +23,10 @@ use crate::test_utils::contracts::FeatureContract;
 use crate::test_utils::declare::declare_tx;
 use crate::test_utils::initial_test_state::test_state;
 use crate::test_utils::{
-    create_calldata, CairoVersion, NonceManager, BALANCE, MAX_FEE, MAX_L1_GAS_AMOUNT,
-    MAX_L1_GAS_PRICE, TEST_ERC20_CONTRACT_ADDRESS,
+    create_calldata, create_trivial_calldata, CairoVersion, NonceManager, BALANCE, MAX_FEE,
+    MAX_L1_GAS_AMOUNT, MAX_L1_GAS_PRICE, TEST_ERC20_CONTRACT_ADDRESS,
 };
+use crate::transaction::account_transaction::AccountTransaction;
 use crate::transaction::constants::DEPLOY_CONTRACT_FUNCTION_ENTRY_POINT_NAME;
 use crate::transaction::objects::FeeType;
 use crate::transaction::test_utils::{
@@ -34,6 +35,45 @@ use crate::transaction::test_utils::{
 use crate::transaction::transaction_execution::Transaction;
 use crate::{declare_tx_args, invoke_tx_args, nonce, storage_key};
 
+fn _trivial_calldata_invoke_tx(
+    account_address: ContractAddress,
+    test_contract_address: ContractAddress,
+    nonce: Nonce,
+) -> AccountTransaction {
+    account_invoke_tx(invoke_tx_args! {
+        sender_address: account_address,
+        calldata: create_trivial_calldata(test_contract_address),
+        resource_bounds: l1_resource_bounds(MAX_L1_GAS_AMOUNT, MAX_L1_GAS_PRICE),
+        version: TransactionVersion::THREE,
+        nonce: nonce,
+    })
+}
+
+/// Checks that the storage values of the account and sequencer balances in the
+/// versioned state of tx_index equals the expected values.
+fn _verify_sequencer_balance_update<S: StateReader>(
+    executor: &WorkerExecutor<'_, S>,
+    tx_index: usize,
+    // We assume the balance is at most 2^128, so the "low" value is sufficient.
+    expected_sequencer_balance_low: StarkFelt,
+) {
+    let tx_version_state = executor.state.pin_version(tx_index);
+    let (sequencer_balance_key_low, sequencer_balance_key_high) =
+        get_sequencer_balance_keys(executor.block_context);
+    for (expected_balance, storage_key) in [
+        (expected_sequencer_balance_low, sequencer_balance_key_low),
+        (StarkFelt::ZERO, sequencer_balance_key_high),
+    ] {
+        let actual_balance = tx_version_state
+            .get_storage_at(
+                executor.block_context.chain_info.fee_token_address(&FeeType::Strk),
+                storage_key,
+            )
+            .unwrap();
+        assert_eq!(expected_balance, actual_balance);
+    }
+}
+
 #[test]
 fn test_worker_execute() {
     // Settings.
diff --git a/crates/blockifier/src/fee/fee_utils.rs b/crates/blockifier/src/fee/fee_utils.rs
index 74f0be60c2..b4626d7010 100644
--- a/crates/blockifier/src/fee/fee_utils.rs
+++ b/crates/blockifier/src/fee/fee_utils.rs
@@ -1,6 +1,7 @@
 use std::collections::HashSet;
 
 use cairo_vm::vm::runners::cairo_runner::ExecutionResources;
+use starknet_api::core::ContractAddress;
 use starknet_api::hash::StarkFelt;
 use starknet_api::state::StorageKey;
 use starknet_api::transaction::Fee;
@@ -147,8 +148,13 @@ pub fn verify_can_pay_committed_bounds(
 
 pub fn get_sequencer_balance_keys(block_context: &BlockContext) -> (StorageKey, StorageKey) {
     let sequencer_address = block_context.block_info.sequencer_address;
-    let sequencer_balance_key_low = get_fee_token_var_address(sequencer_address);
-    let sequencer_balance_key_high = next_storage_key(&sequencer_balance_key_low)
-        .expect("Cannot get sequencer balance high key.");
-    (sequencer_balance_key_low, sequencer_balance_key_high)
+    get_address_balance_keys(sequencer_address)
+}
+
+pub fn get_address_balance_keys(address: ContractAddress) -> (StorageKey, StorageKey) {
+    let balance_key_low = get_fee_token_var_address(address);
+    let balance_key_high = next_storage_key(&balance_key_low).unwrap_or_else(|_| {
+        panic!("Failed to get balance_key_high for address: {:?}", address.0);
+    });
+    (balance_key_low, balance_key_high)
 }