refactor(starknet_integration_tests): add TestScenario to enable test…

…ing different scenarios
starkware-libs · Feb 2, 2025 · e58d457 · e58d457
1 parent 72d83ce
commit e58d457
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 10 deletions.
diff --git a/crates/starknet_integration_tests/src/end_to_end_integration.rs b/crates/starknet_integration_tests/src/end_to_end_integration.rs
@@ -5,7 +5,7 @@ use starknet_api::block::BlockNumber;
 use starknet_sequencer_node::test_utils::node_runner::get_node_executable_path;
 use tracing::info;
 
-use crate::sequencer_manager::{get_sequencer_setup_configs, IntegrationTestManager};
+use crate::{sequencer_manager::{get_sequencer_setup_configs, IntegrationTestManager}, utils::InvokeTxs};
 
 pub async fn end_to_end_integration(tx_generator: &mut MultiAccountTransactionGenerator) {
     const EXPECTED_BLOCK_NUMBER: BlockNumber = BlockNumber(15);
@@ -32,7 +32,13 @@ pub async fn end_to_end_integration(tx_generator: &mut MultiAccountTransactionGe
 
     // Run the test.
     integration_test_manager
-        .test_and_verify(tx_generator, 0, N_TXS, SENDER_ACCOUNT, EXPECTED_BLOCK_NUMBER)
+        .test_and_verify(
+            tx_generator,
+            0,
+            InvokeTxs(N_TXS),
+            SENDER_ACCOUNT,
+            EXPECTED_BLOCK_NUMBER,
+        )
         .await;
 
     // Run the late node.
@@ -43,7 +49,7 @@ pub async fn end_to_end_integration(tx_generator: &mut MultiAccountTransactionGe
         .test_and_verify(
             tx_generator,
             N_TXS,
-            N_TXS,
+            InvokeTxs(N_TXS),
             SENDER_ACCOUNT,
             LATE_NODE_EXPECTED_BLOCK_NUMBER,
         )

diff --git a/crates/starknet_integration_tests/src/sequencer_manager.rs b/crates/starknet_integration_tests/src/sequencer_manager.rs
@@ -37,6 +37,7 @@ use crate::utils::{
     create_mempool_p2p_configs,
     create_state_sync_configs,
     send_account_txs,
+    TestScenario,
 };
 
 /// The number of consolidated local sequencers that participate in the test.
@@ -223,7 +224,7 @@ impl IntegrationTestManager {
         &mut self,
         tx_generator: &mut MultiAccountTransactionGenerator,
         expected_initial_value: usize,
-        n_txs: usize,
+        test_scenario: impl TestScenario,
         sender_account: AccountId,
         expected_block_number: BlockNumber,
     ) {
@@ -233,11 +234,11 @@ impl IntegrationTestManager {
             expected_initial_value + 1,
         )
         .await;
-        self.run_integration_test_simulator(tx_generator, n_txs, sender_account).await;
+        self.run_integration_test_simulator(tx_generator, &test_scenario, sender_account).await;
         self.await_execution(expected_block_number).await;
         self.verify_results(
             tx_generator.account_with_id(sender_account).sender_address(),
-            expected_initial_value + n_txs + 1,
+            expected_initial_value + test_scenario.n_txs() + 1,
         )
         .await;
     }
@@ -280,14 +281,16 @@ impl IntegrationTestManager {
     pub async fn run_integration_test_simulator(
         &self,
         tx_generator: &mut MultiAccountTransactionGenerator,
-        n_txs: usize,
+        test_scenario: &impl TestScenario,
         sender_account: AccountId,
     ) {
         info!("Running integration test simulator.");
         let send_rpc_tx_fn = &mut |rpc_tx| self.send_rpc_tx_fn(rpc_tx);
 
+        let n_txs = test_scenario.n_txs();
         info!("Sending {n_txs} txs.");
-        let tx_hashes = send_account_txs(tx_generator, sender_account, n_txs, send_rpc_tx_fn).await;
+        let tx_hashes =
+            send_account_txs(tx_generator, sender_account, test_scenario, send_rpc_tx_fn).await;
         assert_eq!(tx_hashes.len(), n_txs);
     }
 

diff --git a/crates/starknet_integration_tests/src/utils.rs b/crates/starknet_integration_tests/src/utils.rs
@@ -59,6 +59,57 @@ pub const UNDEPLOYED_ACCOUNT_ID: AccountId = 2;
 // Transactions per second sent to the gateway. This rate makes each block contain ~10 transactions
 // with the set [TimeoutsConfig] .
 pub const TPS: u64 = 2;
+const N_TXS_IN_FIRST_BLOCK: usize = 2;
+
+pub trait TestScenario {
+    fn create_txs(
+        &self,
+        tx_generator: &mut MultiAccountTransactionGenerator,
+        account_id: AccountId,
+    ) -> Vec<RpcTransaction>;
+
+    fn n_txs(&self) -> usize;
+}
+
+pub struct InvokeTxs(pub usize);
+
+impl TestScenario for InvokeTxs {
+    fn create_txs(
+        &self,
+        tx_generator: &mut MultiAccountTransactionGenerator,
+        account_id: AccountId,
+    ) -> Vec<RpcTransaction> {
+        create_invoke_txs(tx_generator, account_id, self.0)
+    }
+
+    fn n_txs(&self) -> usize {
+        self.0
+    }
+}
+
+pub struct FirstBlock;
+
+impl TestScenario for FirstBlock {
+    fn create_txs(
+        &self,
+        tx_generator: &mut MultiAccountTransactionGenerator,
+        account_id: AccountId,
+    ) -> Vec<RpcTransaction> {
+        let txs = create_deploy_account_tx_and_invoke_tx(tx_generator, account_id);
+        assert_eq!(
+            txs.len(),
+            N_TXS_IN_FIRST_BLOCK,
+            "First block should contain exactly {} transactions, but {} transactions were created",
+            N_TXS_IN_FIRST_BLOCK,
+            txs.len(),
+        );
+        txs
+    }
+
+    fn n_txs(&self) -> usize {
+        N_TXS_IN_FIRST_BLOCK
+    }
+}
 
 pub fn create_chain_info() -> ChainInfo {
     let mut chain_info = ChainInfo::create_for_testing();
@@ -333,13 +384,13 @@ pub fn test_many_invoke_txs(tx_hashes: &[TransactionHash]) -> Vec<TransactionHas
 pub async fn send_account_txs<'a, Fut>(
     tx_generator: &mut MultiAccountTransactionGenerator,
     account_id: AccountId,
-    n_txs: usize,
+    test_scenario: &impl TestScenario,
     send_rpc_tx_fn: &'a mut dyn FnMut(RpcTransaction) -> Fut,
 ) -> Vec<TransactionHash>
 where
     Fut: Future<Output = TransactionHash> + 'a,
 {
-    let rpc_txs = create_invoke_txs(tx_generator, account_id, n_txs);
+    let rpc_txs = test_scenario.create_txs(tx_generator, account_id);
     send_rpc_txs(rpc_txs, send_rpc_tx_fn).await
 }