Adjust tests with AccountFetcher and async new()

Author: snawaz

magicblock-committor-service/src/tasks/account_fetcher.rs

@@ -5,6 +5,9 @@ use solana_pubkey::Pubkey;
 use solana_rpc_client::nonblocking::rpc_client::RpcClient;
 use solana_sdk::commitment_config::CommitmentConfig;
 
+//
+// AccountFetcher is used by CommitTask
+//
 pub struct AccountFetcher {
     rpc_client: RpcClient,
 }

magicblock-committor-service/src/tasks/mod.rs

@@ -24,7 +24,7 @@ use thiserror::Error;
 
 use crate::tasks::visitor::Visitor;
 
-use self::account_fetcher::AccountFetcher;
+use account_fetcher::*;
 
 pub mod args_task;
 pub mod buffer_task;
@@ -34,7 +34,7 @@ pub(crate) mod task_visitors;
 pub mod utils;
 pub mod visitor;
 
-mod account_fetcher;
+pub mod account_fetcher;
 
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub enum TaskType {
@@ -470,51 +470,57 @@ mod serialization_safety_test {
     }
 
     // Test BufferTask variants
-    #[test]
-    fn test_buffer_task_instruction_serialization() {
+    #[tokio::test]
+    async fn test_buffer_task_instruction_serialization() {
         let validator = Pubkey::new_unique();
 
-        let buffer_task = BufferTask::new_preparation_required(
-            BufferTaskType::Commit(CommitTask::new(
-                456,
-                false,
-                CommittedAccount {
-                    pubkey: Pubkey::new_unique(),
-                    account: Account {
-                        lamports: 2000,
-                        data: vec![7, 8, 9],
-                        owner: Pubkey::new_unique(),
-                        executable: false,
-                        rent_epoch: 0,
+        let buffer_task =
+            BufferTask::new_preparation_required(BufferTaskType::Commit(
+                CommitTask::new(
+                    456,
+                    false,
+                    CommittedAccount {
+                        pubkey: Pubkey::new_unique(),
+                        account: Account {
+                            lamports: 2000,
+                            data: vec![7, 8, 9],
+                            owner: Pubkey::new_unique(),
+                            executable: false,
+                            rent_epoch: 0,
+                        },
                     },
-                },
-            )),
-        );
+                    AccountFetcher::new(),
+                )
+                .await,
+            ));
         assert_serializable(&buffer_task.instruction(&validator));
     }
 
     // Test preparation instructions
-    #[test]
-    fn test_preparation_instructions_serialization() {
+    #[tokio::test]
+    async fn test_preparation_instructions_serialization() {
         let authority = Pubkey::new_unique();
 
         // Test BufferTask preparation
-        let buffer_task = BufferTask::new_preparation_required(
-            BufferTaskType::Commit(CommitTask::new(
-                789,
-                true,
-                CommittedAccount {
-                    pubkey: Pubkey::new_unique(),
-                    account: Account {
-                        lamports: 3000,
-                        data: vec![0; 1024], // Larger data to test chunking
-                        owner: Pubkey::new_unique(),
-                        executable: false,
-                        rent_epoch: 0,
+        let buffer_task =
+            BufferTask::new_preparation_required(BufferTaskType::Commit(
+                CommitTask::new(
+                    789,
+                    true,
+                    CommittedAccount {
+                        pubkey: Pubkey::new_unique(),
+                        account: Account {
+                            lamports: 3000,
+                            data: vec![0; 1024], // Larger data to test chunking
+                            owner: Pubkey::new_unique(),
+                            executable: false,
+                            rent_epoch: 0,
+                        },
                     },
-                },
-            )),
-        );
+                    AccountFetcher::new(),
+                )
+                .await,
+            ));
 
         let PreparationState::Required(preparation_task) =
             buffer_task.preparation_state()

test-integration/test-committor-service/tests/common.rs

@@ -14,7 +14,7 @@ use magicblock_committor_service::{
         },
         IntentExecutorImpl,
     },
-    tasks::CommitTask,
+    tasks::{account_fetcher::AccountFetcher, CommitTask},
     transaction_preparator::{
         delivery_preparator::DeliveryPreparator, TransactionPreparatorImpl,
     },
@@ -147,7 +147,7 @@ pub fn generate_random_bytes(length: usize) -> Vec<u8> {
 }
 
 #[allow(dead_code)]
-pub fn create_commit_task(data: &[u8]) -> CommitTask {
+pub async fn create_commit_task(data: &[u8]) -> CommitTask {
     static COMMIT_ID: AtomicU64 = AtomicU64::new(0);
     CommitTask::new(
         COMMIT_ID.fetch_add(1, Ordering::Relaxed),
@@ -162,7 +162,9 @@ pub fn create_commit_task(data: &[u8]) -> CommitTask {
                 rent_epoch: 0,
             },
         },
+        AccountFetcher::new(),
     )
+    .await
 }
 
 #[allow(dead_code)]

test-integration/test-committor-service/tests/test_delivery_preparator.rs

@@ -1,4 +1,5 @@
 use borsh::BorshDeserialize;
+use futures::future::join_all;
 use magicblock_committor_program::Chunks;
 use magicblock_committor_service::{
     persist::IntentPersisterImpl,
@@ -23,7 +24,7 @@ async fn test_prepare_10kb_buffer() {
     let preparator = fixture.create_delivery_preparator();
 
     let data = generate_random_bytes(10 * 1024);
-    let buffer_task = BufferTaskType::Commit(create_commit_task(&data));
+    let buffer_task = BufferTaskType::Commit(create_commit_task(&data).await);
     let mut strategy = TransactionStrategy {
         optimized_tasks: vec![Box::new(BufferTask::new_preparation_required(
             buffer_task,
@@ -90,15 +91,13 @@ async fn test_prepare_multiple_buffers() {
         generate_random_bytes(10),
         generate_random_bytes(500 * 1024),
     ];
-    let buffer_tasks = datas
-        .iter()
-        .map(|data| {
-            let task =
-                BufferTaskType::Commit(create_commit_task(data.as_slice()));
-            Box::new(BufferTask::new_preparation_required(task))
-                as Box<dyn BaseTask>
-        })
-        .collect();
+    let buffer_tasks = join_all(datas.iter().map(|data| async {
+        let task =
+            BufferTaskType::Commit(create_commit_task(data.as_slice()).await);
+        Box::new(BufferTask::new_preparation_required(task))
+            as Box<dyn BaseTask>
+    }))
+    .await;
     let mut strategy = TransactionStrategy {
         optimized_tasks: buffer_tasks,
         lookup_tables_keys: vec![],
@@ -169,14 +168,12 @@ async fn test_lookup_tables() {
         generate_random_bytes(20),
         generate_random_bytes(30),
     ];
-    let tasks = datas
-        .iter()
-        .map(|data| {
-            let task =
-                ArgsTaskType::Commit(create_commit_task(data.as_slice()));
-            Box::<ArgsTask>::new(task.into()) as Box<dyn BaseTask>
-        })
-        .collect::<Vec<_>>();
+    let tasks = join_all(datas.iter().map(|data| async {
+        let task =
+            ArgsTaskType::Commit(create_commit_task(data.as_slice()).await);
+        Box::<ArgsTask>::new(task.into()) as Box<dyn BaseTask>
+    }))
+    .await;
 
     let lookup_tables_keys = TaskStrategist::collect_lookup_table_keys(
         &fixture.authority.pubkey(),
@@ -218,7 +215,7 @@ async fn test_already_initialized_error_handled() {
     let preparator = fixture.create_delivery_preparator();
 
     let data = generate_random_bytes(10 * 1024);
-    let mut task = create_commit_task(&data);
+    let mut task = create_commit_task(&data).await;
     let buffer_task = BufferTaskType::Commit(task.clone());
     let mut strategy = TransactionStrategy {
         optimized_tasks: vec![Box::new(BufferTask::new_preparation_required(
@@ -309,9 +306,9 @@ async fn test_prepare_cleanup_and_reprepare_mixed_tasks() {
     let buf_b_data = generate_random_bytes(64 * 1024 + 3);
 
     // Keep these around to modify data later (same commit IDs, different data)
-    let mut commit_args = create_commit_task(&args_data);
-    let mut commit_a = create_commit_task(&buf_a_data);
-    let mut commit_b = create_commit_task(&buf_b_data);
+    let mut commit_args = create_commit_task(&args_data).await;
+    let mut commit_a = create_commit_task(&buf_a_data).await;
+    let mut commit_b = create_commit_task(&buf_b_data).await;
 
     let mut strategy = TransactionStrategy {
         optimized_tasks: vec![

test-integration/test-committor-service/tests/test_transaction_preparator.rs

@@ -3,6 +3,7 @@ use magicblock_committor_program::Chunks;
 use magicblock_committor_service::{
     persist::IntentPersisterImpl,
     tasks::{
+        account_fetcher::AccountFetcher,
         args_task::{ArgsTask, ArgsTaskType},
         buffer_task::{BufferTask, BufferTaskType},
         task_strategist::{TaskStrategist, TransactionStrategy},
@@ -35,11 +36,15 @@ async fn test_prepare_commit_tx_with_single_account() {
     let committed_account = create_committed_account(&account_data);
 
     let tasks = vec![
-        Box::new(ArgsTask::new(ArgsTaskType::Commit(CommitTask::new(
-            1,
-            true,
-            committed_account.clone(),
-        )))) as Box<dyn BaseTask>,
+        Box::new(ArgsTask::new(ArgsTaskType::Commit(
+            CommitTask::new(
+                1,
+                true,
+                committed_account.clone(),
+                AccountFetcher::new(),
+            )
+            .await,
+        ))) as Box<dyn BaseTask>,
         Box::new(ArgsTask::new(ArgsTaskType::Finalize(FinalizeTask {
             delegated_account: committed_account.pubkey,
         }))),
@@ -91,16 +96,26 @@ async fn test_prepare_commit_tx_with_multiple_accounts() {
 
     let buffer_commit_task =
         BufferTask::new_preparation_required(BufferTaskType::Commit(
-            CommitTask::new(1, true, committed_account2.clone()),
+            CommitTask::new(
+                1,
+                true,
+                committed_account2.clone(),
+                AccountFetcher::new(),
+            )
+            .await,
         ));
     // Create test data
     let tasks = vec![
         // account 1
-        Box::new(ArgsTask::new(ArgsTaskType::Commit(CommitTask::new(
-            1,
-            true,
-            committed_account1.clone(),
-        )))) as Box<dyn BaseTask>,
+        Box::new(ArgsTask::new(ArgsTaskType::Commit(
+            CommitTask::new(
+                1,
+                true,
+                committed_account1.clone(),
+                AccountFetcher::new(),
+            )
+            .await,
+        ))) as Box<dyn BaseTask>,
         // account 2
         Box::new(buffer_commit_task),
         // finalize account 1
@@ -185,7 +200,13 @@ async fn test_prepare_commit_tx_with_base_actions() {
 
     let buffer_commit_task =
         BufferTask::new_preparation_required(BufferTaskType::Commit(
-            CommitTask::new(1, true, committed_account.clone()),
+            CommitTask::new(
+                1,
+                true,
+                committed_account.clone(),
+                AccountFetcher::new(),
+            )
+            .await,
         ));
     let tasks = vec![
         // commit account