fix tests

Author: snawaz

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

@@ -25,13 +25,10 @@ use magicblock_rpc_client::{
     MagicBlockRpcClientError, MagicBlockSendTransactionConfig,
     MagicBlockSendTransactionOutcome, MagicblockRpcClient,
 };
-
-use solana_keypair::Keypair;
-use solana_message::VersionedMessage;
-
 #[cfg(any(test, feature = "dev-context-only-utils"))]
 pub use null_task_info_fetcher::*;
-
+use solana_keypair::Keypair;
+use solana_message::VersionedMessage;
 use solana_pubkey::Pubkey;
 use solana_rpc_client_api::config::RpcTransactionConfig;
 use solana_signature::Signature;

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

@@ -4,6 +4,8 @@ use magicblock_metrics::metrics::LabelValue;
 use solana_instruction::Instruction;
 use solana_pubkey::Pubkey;
 
+#[cfg(any(test, feature = "dev-context-only-utils"))]
+use super::args_task::ArgsTaskType;
 #[cfg(test)]
 use crate::tasks::TaskStrategy;
 use crate::{
@@ -62,6 +64,17 @@ impl BufferTask {
     }
 }
 
+#[cfg(any(test, feature = "dev-context-only-utils"))]
+impl From<ArgsTaskType> for BufferTaskType {
+    fn from(value: ArgsTaskType) -> Self {
+        match value {
+            ArgsTaskType::Commit(task) => BufferTaskType::Commit(task),
+            ArgsTaskType::CommitDiff(_) => panic!("BufferTask doesn't support CommitDiff yet. Disable your tests temporarily till the next PR"),
+            _ => unimplemented!("Only commit task can be BufferTask currently. Fix your tests"),
+        }
+    }
+}
+
 impl BaseTask for BufferTask {
     fn instruction(&self, validator: &Pubkey) -> Instruction {
         let BufferTaskType::Commit(ref value) = self.task_type;

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

@@ -29,6 +29,8 @@ pub(crate) mod task_visitors;
 pub mod utils;
 pub mod visitor;
 
+pub use task_builder::TaskBuilderImpl;
+
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub enum TaskType {
     Commit,
@@ -301,6 +303,8 @@ pub type BaseTaskResult<T> = Result<T, BaseTaskError>;
 
 #[cfg(test)]
 mod serialization_safety_test {
+    use std::sync::Arc;
+
     use magicblock_program::{
         args::ShortAccountMeta, magic_scheduled_base_intent::ProgramArgs,
     };
@@ -310,7 +314,7 @@ mod serialization_safety_test {
         intent_executor::NullTaskInfoFetcher,
         tasks::{
             args_task::{ArgsTask, ArgsTaskType},
-            buffer_task::{BufferTask, BufferTaskType},
+            buffer_task::BufferTask,
             task_builder::TaskBuilderImpl,
             *,
         },
@@ -382,25 +386,26 @@ mod serialization_safety_test {
     async fn test_buffer_task_instruction_serialization() {
         let validator = Pubkey::new_unique();
 
-        let buffer_task =
-            BufferTask::new_preparation_required(BufferTaskType::Commit(
-                TaskBuilderImpl::create_commit_task(
-                    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(
+            TaskBuilderImpl::create_commit_task(
+                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,
                     },
-                    &Arc::new(NullTaskInfoFetcher),
-                )
-                .await,
-            ));
+                },
+                &Arc::new(NullTaskInfoFetcher),
+            )
+            .await
+            .task_type
+            .into(),
+        );
         assert_serializable(&buffer_task.instruction(&validator));
     }
 
@@ -410,25 +415,26 @@ mod serialization_safety_test {
         let authority = Pubkey::new_unique();
 
         // Test BufferTask preparation
-        let buffer_task =
-            BufferTask::new_preparation_required(BufferTaskType::Commit(
-                TaskBuilderImpl::create_commit_task(
-                    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(
+            TaskBuilderImpl::create_commit_task(
+                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,
                     },
-                    &Arc::new(NullTaskInfoFetcher),
-                )
-                .await,
-            ));
+                },
+                &Arc::new(NullTaskInfoFetcher),
+            )
+            .await
+            .task_type
+            .into(),
+        );
 
         let PreparationState::Required(preparation_task) =
             buffer_task.preparation_state()

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

@@ -10,6 +10,7 @@ use magicblock_program::magic_scheduled_base_intent::{
 use solana_pubkey::Pubkey;
 use solana_signature::Signature;
 
+use super::{CommitDiffTask, CommitTask};
 use crate::{
     intent_executor::task_info_fetcher::{
         TaskInfoFetcher, TaskInfoFetcherError,
@@ -21,8 +22,6 @@ use crate::{
     },
 };
 
-use super::{CommitDiffTask, CommitTask};
-
 #[async_trait]
 pub trait TasksBuilder {
     // Creates tasks for commit stage
@@ -43,10 +42,12 @@ pub trait TasksBuilder {
 /// V1: Actions are part of finalize tx
 pub struct TaskBuilderImpl;
 
-// Accounts larger than COMMIT_STATE_SIZE_THRESHOLD, use CommitDiff to
-// reduce instruction size. Below this, commit is sent as CommitState.
-// Chose 256 as thresold seems good enough as it could hold 8 u32 fields
-// or 4 u64 fields.
+// Accounts larger than COMMIT_STATE_SIZE_THRESHOLD use CommitDiff to
+// reduce instruction size. Below this threshold, the commit is sent
+// as CommitState. The value (256) is chosen because it is sufficient
+// for small accounts, which typically could hold up to 8 u32 fields or
+// 4 u64 fields. These integers are expected to be on the hot path
+// and updated continuously.
 const COMMIT_STATE_SIZE_THRESHOLD: usize = 256;
 
 impl TaskBuilderImpl {
@@ -78,14 +79,14 @@ impl TaskBuilderImpl {
 
         if let Some(base_account) = base_account {
             ArgsTaskType::CommitDiff(CommitDiffTask {
-                commit_id: commit_id,
+                commit_id,
                 allow_undelegation,
                 committed_account: account,
                 base_account,
             })
         } else {
             ArgsTaskType::Commit(CommitTask {
-                commit_id: commit_id,
+                commit_id,
                 allow_undelegation,
                 committed_account: account,
             })

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

@@ -373,27 +373,30 @@ mod tests {
     use magicblock_program::magic_scheduled_base_intent::{
         BaseAction, CommittedAccount, ProgramArgs,
     };
+    use magicblock_rpc_client::MagicBlockRpcClientResult;
     use solana_account::Account;
+    use solana_program::system_program;
     use solana_pubkey::Pubkey;
     use solana_system_program::id as system_program_id;
 
     use super::*;
     use crate::{
         intent_execution_manager::intent_scheduler::create_test_intent,
-        intent_executor::task_info_fetcher::NullTaskInfoFetcher,
-        intent_executor::task_info_fetcher::{
-            ResetType, TaskInfoFetcher, TaskInfoFetcherResult,
+        intent_executor::{
+            task_info_fetcher::{
+                ResetType, TaskInfoFetcher, TaskInfoFetcherResult,
+            },
+            NullTaskInfoFetcher,
         },
         persist::IntentPersisterImpl,
         tasks::{
-            intent_executor::NullTaskInfoFetcher,
-            persist::IntentPersisterImpl,
             task_builder::{TaskBuilderImpl, TasksBuilder},
-            BaseActionTask, CommitTask, TaskStrategy, UndelegateTask,
+            BaseActionTask, TaskStrategy, UndelegateTask,
         },
     };
 
     struct MockInfoFetcher;
+
     #[async_trait::async_trait]
     impl TaskInfoFetcher for MockInfoFetcher {
         async fn fetch_next_commit_ids(
@@ -415,31 +418,36 @@ mod tests {
         }
 
         fn reset(&self, _: ResetType) {}
+
+        async fn get_base_account(
+            &self,
+            _pubkey: &Pubkey,
+        ) -> MagicBlockRpcClientResult<Option<Account>> {
+            Ok(None) // Account Not Found
+        }
     }
 
     // Helper to create a simple commit task
     async fn create_test_commit_task(
         commit_id: u64,
         data_size: usize,
     ) -> ArgsTask {
-        ArgsTask::new(ArgsTaskType::Commit(
-            TaskBuilderImpl::create_commit_task(
-                commit_id,
-                false,
-                CommittedAccount {
-                    pubkey: Pubkey::new_unique(),
-                    account: Account {
-                        lamports: 1000,
-                        data: vec![1; data_size],
-                        owner: system_program::id(),
-                        executable: false,
-                        rent_epoch: 0,
-                    },
+        TaskBuilderImpl::create_commit_task(
+            commit_id,
+            false,
+            CommittedAccount {
+                pubkey: Pubkey::new_unique(),
+                account: Account {
+                    lamports: 1000,
+                    data: vec![1; data_size],
+                    owner: system_program::id(),
+                    executable: false,
+                    rent_epoch: 0,
                 },
-                &Arc::new(NullTaskInfoFetcher),
-            )
-            .await,
-        ))
+            },
+            &Arc::new(NullTaskInfoFetcher),
+        )
+        .await
     }
 
     // Helper to create a Base action task

test-integration/programs/schedulecommit/src/lib.rs

@@ -15,8 +15,7 @@ use solana_program::{
     entrypoint::{self, ProgramResult},
     instruction::{AccountMeta, Instruction},
     msg,
-    program::invoke,
-    program::invoke_signed,
+    program::{invoke, invoke_signed},
     program_error::ProgramError,
     pubkey::Pubkey,
     rent::Rent,

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

@@ -12,9 +12,9 @@ use magicblock_committor_service::{
         task_info_fetcher::{
             ResetType, TaskInfoFetcher, TaskInfoFetcherResult,
         },
-        IntentExecutorImpl, NullTaskInfoFetcher,
+        IntentExecutorImpl,
     },
-    tasks::{CommitTask, TaskBuilderImpl},
+    tasks::CommitTask,
     transaction_preparator::{
         delivery_preparator::DeliveryPreparator, TransactionPreparatorImpl,
     },
@@ -159,12 +159,12 @@ pub fn generate_random_bytes(length: usize) -> Vec<u8> {
 }
 
 #[allow(dead_code)]
-pub async fn create_commit_task(data: &[u8]) -> CommitTask {
+pub fn create_commit_task(data: &[u8]) -> CommitTask {
     static COMMIT_ID: AtomicU64 = AtomicU64::new(0);
-    TaskBuilderImpl::create_commit_task(
-        COMMIT_ID.fetch_add(1, Ordering::Relaxed),
-        false,
-        CommittedAccount {
+    CommitTask {
+        commit_id: COMMIT_ID.fetch_add(1, Ordering::Relaxed),
+        allow_undelegation: false,
+        committed_account: CommittedAccount {
             pubkey: Pubkey::new_unique(),
             account: Account {
                 lamports: 1000,
@@ -174,9 +174,7 @@ pub async fn create_commit_task(data: &[u8]) -> CommitTask {
                 rent_epoch: 0,
             },
         },
-        &Arc::new(NullTaskInfoFetcher),
-    )
-    .await
+    }
 }
 
 #[allow(dead_code)]