chore: unit tests for queue filtering function

Author: IAvecilla

batcher/aligned-batcher/src/lib.rs

@@ -47,6 +47,9 @@ pub mod sp1;
 pub mod types;
 mod zk_utils;
 
+#[cfg(test)]
+mod testonly;
+
 const AGGREGATOR_GAS_COST: u128 = 400_000;
 const BATCHER_SUBMISSION_BASE_GAS_COST: u128 = 125_000;
 pub(crate) const ADDITIONAL_SUBMISSION_GAS_COST_PER_PROOF: u128 = 13_000;
@@ -426,7 +429,9 @@ impl Batcher {
                 );
                 send_message(
                     ws_conn_sink.clone(),
-                    ValidityResponseMessage::InvalidProof(ProofInvalidReason::DisabledVerifier),
+                    ValidityResponseMessage::InvalidProof(ProofInvalidReason::DisabledVerifier(
+                        verification_data.proving_system,
+                    )),
                 )
                 .await;
                 return Ok(());

batcher/aligned-batcher/src/testonly/mod.rs

@@ -0,0 +1 @@
+pub mod queue_entry;

batcher/aligned-batcher/src/testonly/queue_entry.rs

@@ -0,0 +1,36 @@
+use aligned_sdk::core::types::{
+    NoncedVerificationData, ProvingSystemId, VerificationData, VerificationDataCommitment,
+};
+use ethers::types::{Address, Signature, U256};
+
+use crate::types::batch_queue::BatchQueueEntry;
+
+impl BatchQueueEntry {
+    pub async fn new_for_test(proving_system: ProvingSystemId, sender: Address) -> Self {
+        let dummy_signature = Signature {
+            r: U256::from(1),
+            s: U256::from(2),
+            v: 3,
+        };
+        BatchQueueEntry {
+            sender,
+            messaging_sink: None,
+            verification_data_commitment: VerificationDataCommitment::default(),
+            nonced_verification_data: NoncedVerificationData {
+                verification_data: VerificationData {
+                    proving_system,
+                    proof: vec![],
+                    pub_input: None,
+                    verification_key: None,
+                    vm_program_code: None,
+                    proof_generator_addr: Address::default(),
+                },
+                nonce: U256::zero(),
+                max_fee: Default::default(),
+                chain_id: Default::default(),
+                payment_service_addr: Default::default(),
+            },
+            signature: dummy_signature,
+        }
+    }
+}

batcher/aligned-batcher/src/zk_utils/mod.rs

@@ -128,9 +128,15 @@ pub(crate) async fn filter_disabled_verifiers(
 
 #[cfg(test)]
 mod test {
+    use crate::{
+        types::batch_queue::{BatchQueue, BatchQueueEntry, BatchQueueEntryPriority},
+        zk_utils::filter_disabled_verifiers,
+    };
+
     use super::is_verifier_disabled;
     use aligned_sdk::core::types::{ProvingSystemId, VerificationData};
-    use ethers::types::Address;
+    use ethers::types::{Address, U256};
+    use tokio::sync::Mutex;
 
     fn get_all_verifiers() -> Vec<ProvingSystemId> {
         let verifiers = vec![
@@ -224,4 +230,124 @@ mod test {
             }
         }
     }
+
+    #[tokio::test]
+    async fn test_remove_disabled_verifiers_from_queue() {
+        // Disabling SP1 verifier.
+        let disabled_verifiers = Mutex::new(U256::from(8));
+
+        let mut batch_queue = BatchQueue::new();
+        let entry_with_sp1_sender_0 =
+            BatchQueueEntry::new_for_test(ProvingSystemId::SP1, Address::zero()).await;
+        let entry_with_risc0_sender_1 =
+            BatchQueueEntry::new_for_test(ProvingSystemId::Risc0, Address::from_slice(&[1; 20]))
+                .await;
+        batch_queue.push(
+            entry_with_sp1_sender_0,
+            BatchQueueEntryPriority::new(0.into(), 0.into()),
+        );
+        batch_queue.push(
+            entry_with_risc0_sender_1,
+            BatchQueueEntryPriority::new(0.into(), 0.into()),
+        );
+
+        assert_eq!(batch_queue.len(), 2);
+        let filtered_queue =
+            filter_disabled_verifiers(batch_queue, disabled_verifiers.lock().await).await;
+        assert_eq!(filtered_queue.len(), 1);
+        let entry = filtered_queue.peek().unwrap().0.clone();
+        assert_eq!(
+            entry
+                .nonced_verification_data
+                .verification_data
+                .proving_system,
+            aligned_sdk::core::types::ProvingSystemId::Risc0
+        );
+    }
+
+    #[tokio::test]
+    async fn test_remove_new_data_from_user_with_disabled_verifier() {
+        // Disabling SP1 verifier.
+        let disabled_verifiers = Mutex::new(U256::from(8));
+        let mut batch_queue = BatchQueue::new();
+
+        let entry_with_sp1_sender_0 =
+            BatchQueueEntry::new_for_test(ProvingSystemId::SP1, Address::zero()).await;
+        let entry_with_risc0_sender_1 =
+            BatchQueueEntry::new_for_test(ProvingSystemId::Risc0, Address::from_slice(&[1; 20]))
+                .await;
+        let mut another_entry_sp1_sender_0 =
+            BatchQueueEntry::new_for_test(ProvingSystemId::SP1, Address::zero()).await;
+        another_entry_sp1_sender_0.nonced_verification_data.nonce = U256::from(1);
+
+        batch_queue.push(
+            entry_with_sp1_sender_0,
+            BatchQueueEntryPriority::new(0.into(), 0.into()),
+        );
+        batch_queue.push(
+            entry_with_risc0_sender_1,
+            BatchQueueEntryPriority::new(0.into(), 0.into()),
+        );
+        batch_queue.push(
+            another_entry_sp1_sender_0,
+            BatchQueueEntryPriority::new(0.into(), 1.into()),
+        );
+        assert_eq!(batch_queue.len(), 3);
+        let filtered_batch_queue =
+            filter_disabled_verifiers(batch_queue.clone(), disabled_verifiers.lock().await).await;
+        assert_eq!(filtered_batch_queue.len(), 1);
+
+        let entry = filtered_batch_queue.peek().unwrap().0.clone();
+        assert_eq!(
+            entry
+                .nonced_verification_data
+                .verification_data
+                .proving_system,
+            aligned_sdk::core::types::ProvingSystemId::Risc0
+        );
+    }
+
+    #[tokio::test]
+    async fn test_keep_old_proofs_for_user_with_new_invalid() {
+        // Disabling SP1 verifier.
+        let disabled_verifiers = Mutex::new(U256::from(8));
+        let mut batch_queue = BatchQueue::new();
+
+        let entry_risc0_sender_0 =
+            BatchQueueEntry::new_for_test(ProvingSystemId::Risc0, Address::zero()).await;
+        let mut entry_sp1_sender_0 =
+            BatchQueueEntry::new_for_test(ProvingSystemId::SP1, Address::zero()).await;
+        entry_sp1_sender_0.nonced_verification_data.nonce = U256::from(1);
+        let mut another_entry_risc0_sender_0 =
+            BatchQueueEntry::new_for_test(ProvingSystemId::Risc0, Address::zero()).await;
+        another_entry_risc0_sender_0.nonced_verification_data.nonce = U256::from(2);
+
+        batch_queue.push(
+            entry_risc0_sender_0,
+            BatchQueueEntryPriority::new(0.into(), 0.into()),
+        );
+        batch_queue.push(
+            entry_sp1_sender_0,
+            BatchQueueEntryPriority::new(0.into(), 1.into()),
+        );
+        batch_queue.push(
+            another_entry_risc0_sender_0,
+            BatchQueueEntryPriority::new(0.into(), 2.into()),
+        );
+        assert_eq!(batch_queue.len(), 3);
+
+        let filtered_batch_queue =
+            filter_disabled_verifiers(batch_queue.clone(), disabled_verifiers.lock().await).await;
+        assert_eq!(filtered_batch_queue.len(), 1);
+
+        let entry = filtered_batch_queue.peek().unwrap().0.clone();
+        assert_eq!(
+            entry
+                .nonced_verification_data
+                .verification_data
+                .proving_system,
+            aligned_sdk::core::types::ProvingSystemId::Risc0
+        );
+        assert_eq!(entry.nonced_verification_data.nonce, U256::zero());
+    }
 }