starknet_os: os resources test - extract setup_test_builder

Author: dorimedini-starkware

crates/starknet_os_flow_tests/src/os_resources_test.rs

@@ -19,7 +19,7 @@ use expect_test::expect_file;
 use indexmap::IndexMap;
 use starknet_api::block::StarknetVersion;
 use starknet_api::contract_class::SierraVersion;
-use starknet_api::core::EthAddress;
+use starknet_api::core::{ClassHash, ContractAddress, EthAddress};
 use starknet_api::executable_transaction::InvokeTransaction;
 use starknet_api::test_utils::invoke::invoke_tx;
 use starknet_api::transaction::fields::ContractAddressSalt;
@@ -107,6 +107,13 @@ const FEE_TRANSFER_SYSCALLS: [Selector; 10] = [
 /// All other syscalls are called only once.
 const SYSCALLS_CALLED_TWICE: [Selector; 2] = [Selector::Secp256k1New, Selector::Secp256r1New];
 
+struct OsResourcesTestSetup {
+    os_resources_contract_address: ContractAddress,
+    stable_contract_address: ContractAddress,
+    stable_contract_class_hash: ClassHash,
+    test_builder: TestBuilder<DictStateReader>,
+}
+
 /// See [SYSCALLS_WITH_VIRTUAL_BUILTINS] for why this function is needed.
 fn update_resources_for_virtual_builtin_syscall(
     selector: Selector,
@@ -133,16 +140,89 @@ fn update_resources_for_virtual_builtin_syscall(
     }
 }
 
+/// Setup an test builder with
+/// 1. the OS-resources contract deployed,
+/// 2. funded (it is an account contract as well),
+/// 3. the initial block context (and therefore subsequent block contexts) with the minimal sierra
+///    version for gas tracking set to "infinity", in order to force step tracking mode, and
+/// 4. the stable contract declared and deployed.
+/// Pass the raw VC used in tests, to make sure the test initial state is set up consistently.
+async fn setup_test_builder(raw_vc: Option<&RawVersionedConstants>) -> OsResourcesTestSetup {
+    // Setup the test initial state and test builder.
+    // Need to explicitly set up the state to be able to override the minimal sierra version for gas
+    // tracking, in order to force step tracking mode.
+    let os_resources_contract = FeatureContract::OsResourcesTest(RunnableCairo1::Casm);
+    let (mut initial_state_data, [os_resources_contract_address]) =
+        create_default_initial_state_data::<DictStateReader, 1>([(
+            os_resources_contract,
+            calldata![],
+        )])
+        .await;
+    initial_state_data.initial_state.block_context = {
+        let block_context = &initial_state_data.initial_state.block_context;
+        let mut vc = block_context.versioned_constants().clone();
+        if let Some(raw_vc) = raw_vc {
+            assert_eq!(vc, raw_vc.clone().into());
+        }
+        vc.min_sierra_version_for_sierra_gas = SierraVersion::new(99, 99, 99);
+        BlockContext::new(
+            block_context.block_info().clone(),
+            block_context.chain_info().clone(),
+            vc,
+            block_context.bouncer_config.clone(),
+        )
+    };
+    let virtual_os = false;
+    let mut test_builder = TestBuilder::new_with_initial_state_data(
+        initial_state_data,
+        TestBuilderConfig::default(),
+        virtual_os,
+    );
+    test_builder.add_fund_address_tx_with_default_amount(os_resources_contract_address);
+
+    // Declare and deploy an instance of the stable contract.
+    let stable_contract_sierra = &DEPLOYABLE_FOR_RESOURCE_MEASUREMENT_CONTRACT_SIERRA;
+    let stable_contract_casm = &DEPLOYABLE_FOR_RESOURCE_MEASUREMENT_CONTRACT_CASM;
+    let stable_contract_class_hash = stable_contract_sierra.calculate_class_hash();
+    let extra_declare_args = declare_tx_args! {
+        sender_address: *FUNDED_ACCOUNT_ADDRESS,
+        nonce: test_builder.next_nonce(*FUNDED_ACCOUNT_ADDRESS),
+        resource_bounds: *NON_TRIVIAL_RESOURCE_BOUNDS,
+    };
+    test_builder.add_explicit_cairo1_declare_tx(
+        stable_contract_sierra,
+        (**stable_contract_casm).clone(),
+        extra_declare_args,
+        &test_builder.chain_id(),
+    );
+    let deploy_from_zero = true;
+    let (deploy_tx, stable_contract_address) =
+        get_deploy_contract_tx_and_address_with_salt_and_deployer(
+            stable_contract_class_hash,
+            calldata![Felt::ZERO], // Ctor calldata length.
+            test_builder.next_nonce(*FUNDED_ACCOUNT_ADDRESS),
+            *NON_TRIVIAL_RESOURCE_BOUNDS,
+            ContractAddressSalt::default(),
+            deploy_from_zero,
+        );
+    test_builder.add_invoke_tx(deploy_tx, None, None);
+
+    // Move on to the next block, so the measurement txs are in their own block.
+    test_builder.move_to_next_block();
+
+    OsResourcesTestSetup {
+        os_resources_contract_address,
+        stable_contract_address,
+        stable_contract_class_hash,
+        test_builder,
+    }
+}
+
 /// Regression test for the list of syscalls called during the fee transfer phase of a transaction.
 #[tokio::test]
 async fn test_fee_transfer_syscalls() {
-    let os_resources_contract = FeatureContract::OsResourcesTest(RunnableCairo1::Casm);
-    let (mut builder, [os_resources_contract_address]) =
-        TestBuilder::create_standard([(os_resources_contract, calldata![])]).await;
-
-    // Fund the contract - it will be used as the account.
-    // Then, move on to the next block, so the syscall-measurement tx is in it's own block.
-    builder.add_fund_address_tx_with_default_amount(os_resources_contract_address);
+    let OsResourcesTestSetup { os_resources_contract_address, test_builder: mut builder, .. } =
+        setup_test_builder(None).await;
 
     // Invoke from the OS resources contract, with zeros as calldata, to make the __execute__ do
     // nothing. All resulting events should be from the fee transfer call.
@@ -201,71 +281,16 @@ async fn test_fee_transfer_syscalls() {
 ///    changes in the deploying contract address are not reflected in the measurements.
 #[tokio::test]
 async fn test_os_resources_regression() {
-    let os_resources_contract = FeatureContract::OsResourcesTest(RunnableCairo1::Casm);
     let version = StarknetVersion::LATEST;
     let mut raw_vc: RawVersionedConstants =
         serde_json::from_str(VersionedConstants::json_str(&version).unwrap()).unwrap();
 
-    // Setup the test initial state and test builder.
-    // Need to explicitly set up the state to be able to override the minimal sierra version for gas
-    // tracking, in order to force step tracking mode.
-    let (mut initial_state_data, [os_resources_contract_address]) =
-        create_default_initial_state_data::<DictStateReader, 1>([(
-            os_resources_contract,
-            calldata![],
-        )])
-        .await;
-    initial_state_data.initial_state.block_context = {
-        let block_context = &initial_state_data.initial_state.block_context;
-        let mut vc = block_context.versioned_constants().clone();
-        assert_eq!(vc, raw_vc.clone().into());
-        vc.min_sierra_version_for_sierra_gas = SierraVersion::new(99, 99, 99);
-        BlockContext::new(
-            block_context.block_info().clone(),
-            block_context.chain_info().clone(),
-            vc,
-            block_context.bouncer_config.clone(),
-        )
-    };
-    let virtual_os = false;
-    let mut test_builder = TestBuilder::new_with_initial_state_data(
-        initial_state_data,
-        TestBuilderConfig::default(),
-        virtual_os,
-    );
-
-    // Fund the contract - it will be used as the account.
-    test_builder.add_fund_address_tx_with_default_amount(os_resources_contract_address);
-
-    // Declare and deploy an instance of the stable contract.
-    let stable_contract_sierra = &DEPLOYABLE_FOR_RESOURCE_MEASUREMENT_CONTRACT_SIERRA;
-    let stable_contract_casm = &DEPLOYABLE_FOR_RESOURCE_MEASUREMENT_CONTRACT_CASM;
-    let stable_contract_class_hash = stable_contract_sierra.calculate_class_hash();
-    let extra_declare_args = declare_tx_args! {
-        sender_address: *FUNDED_ACCOUNT_ADDRESS,
-        nonce: test_builder.next_nonce(*FUNDED_ACCOUNT_ADDRESS),
-        resource_bounds: *NON_TRIVIAL_RESOURCE_BOUNDS,
-    };
-    test_builder.add_explicit_cairo1_declare_tx(
-        stable_contract_sierra,
-        (**stable_contract_casm).clone(),
-        extra_declare_args,
-        &test_builder.chain_id(),
-    );
-    let deploy_from_zero = true;
-    let (deploy_tx, stable_contract_address) =
-        get_deploy_contract_tx_and_address_with_salt_and_deployer(
-            stable_contract_class_hash,
-            calldata![Felt::ZERO], // Ctor calldata length.
-            test_builder.next_nonce(*FUNDED_ACCOUNT_ADDRESS),
-            *NON_TRIVIAL_RESOURCE_BOUNDS,
-            ContractAddressSalt::default(),
-            deploy_from_zero,
-        );
-    test_builder.add_invoke_tx(deploy_tx, None, None);
-
-    // Move on to the next block, so the syscall-measurement tx is in it's own block.
-    test_builder.move_to_next_block();
+    let OsResourcesTestSetup {
+        os_resources_contract_address,
+        stable_contract_address,
+        stable_contract_class_hash,
+        mut test_builder,
+    } = setup_test_builder(Some(&raw_vc)).await;
 
     // Add the syscall-measurement tx.
     let tx = InvokeTransaction::create(