apollo_batcher: test fetch revert result

Author: yoavGrs

crates/apollo_batcher/src/commitment_manager/commitment_manager_test.rs

@@ -59,12 +59,13 @@ fn get_number_of_tasks_in_receiver<T>(receiver: &Receiver<T>) -> usize {
 
 async fn create_mock_commitment_manager(
     mock_dependencies: MockDependencies,
+    commitment_manager_config: Option<CommitmentManagerConfig>,
 ) -> MockCommitmentManager {
-    let commitment_manager_config = CommitmentManagerConfig {
+    let commitment_manager_config = commitment_manager_config.unwrap_or(CommitmentManagerConfig {
         tasks_channel_size: 1,
         results_channel_size: 1,
         wait_for_tasks_channel: false,
-    };
+    });
     CommitmentManager::create_commitment_manager(
         &mock_dependencies.batcher_config,
         // TODO(Amos): Use commitment manager config in batcher config, once it's added.
@@ -75,7 +76,7 @@ async fn create_mock_commitment_manager(
     .await
 }
 
-async fn await_results<T>(receiver: &mut Receiver<T>, expected_n_results: usize) -> Vec<T> {
+async fn wait_for_filling_channel<T>(receiver: &mut Receiver<T>, expected_n_results: usize) {
     let max_n_retries = 3;
     let mut n_retries = 0;
     while get_number_of_tasks_in_receiver(receiver) < expected_n_results {
@@ -88,6 +89,10 @@ async fn await_results<T>(receiver: &mut Receiver<T>, expected_n_results: usize)
             );
         }
     }
+}
+
+async fn await_results<T>(receiver: &mut Receiver<T>, expected_n_results: usize) -> Vec<T> {
+    wait_for_filling_channel(receiver, expected_n_results).await;
     let mut results = Vec::new();
     while let Ok(result) = receiver.try_recv() {
         results.push(result);
@@ -104,7 +109,7 @@ async fn await_results<T>(receiver: &mut Receiver<T>, expected_n_results: usize)
 #[tokio::test]
 async fn test_create_commitment_manager(mut mock_dependencies: MockDependencies) {
     add_initial_heights(&mut mock_dependencies);
-    let commitment_manager = create_mock_commitment_manager(mock_dependencies).await;
+    let commitment_manager = create_mock_commitment_manager(mock_dependencies, None).await;
 
     assert_eq!(
         commitment_manager.get_commitment_task_offset(),
@@ -140,7 +145,7 @@ async fn test_create_commitment_manager_with_missing_tasks(
         .with(eq(global_root_height))
         .returning(|_| Ok(Some(test_state_diff())));
 
-    let mut commitment_manager = create_mock_commitment_manager(mock_dependencies).await;
+    let mut commitment_manager = create_mock_commitment_manager(mock_dependencies, None).await;
 
     assert_eq!(commitment_manager.get_commitment_task_offset(), INITIAL_HEIGHT,);
     assert_eq!(get_number_of_tasks_in_sender(&commitment_manager.tasks_sender), 1,);
@@ -157,7 +162,7 @@ async fn test_add_commitment_task(mut mock_dependencies: MockDependencies) {
     let state_diff = test_state_diff();
     let state_diff_commitment = Some(StateDiffCommitment::default());
 
-    let mut commitment_manager = create_mock_commitment_manager(mock_dependencies).await;
+    let mut commitment_manager = create_mock_commitment_manager(mock_dependencies, None).await;
 
     // Verify incorrect height results in error.
     let incorrect_height = INITIAL_HEIGHT.next().unwrap();
@@ -217,7 +222,7 @@ async fn test_add_commitment_task_full(mut mock_dependencies: MockDependencies)
     let state_diff = test_state_diff();
     let state_diff_commitment = Some(StateDiffCommitment::default());
 
-    let mut commitment_manager = create_mock_commitment_manager(mock_dependencies).await;
+    let mut commitment_manager = create_mock_commitment_manager(mock_dependencies, None).await;
 
     assert_eq!(
         commitment_manager.config.tasks_channel_size, 1,
@@ -251,14 +256,8 @@ async fn test_get_commitment_results(mut mock_dependencies: MockDependencies) {
         results_channel_size: 2,
         wait_for_tasks_channel: false,
     };
-    let mut commitment_manager = MockCommitmentManager::create_commitment_manager(
-        &mock_dependencies.batcher_config,
-        // TODO(Amos): Use commitment manager config in batcher config, once it's added.
-        &commitment_manager_config,
-        &mock_dependencies.storage_reader,
-        Arc::new(mock_dependencies.committer_client),
-    )
-    .await;
+    let mut commitment_manager =
+        create_mock_commitment_manager(mock_dependencies, Some(commitment_manager_config)).await;
 
     // Verify the commitment manager doesn't wait if there are no results.
     let results = commitment_manager.get_commitment_results().await;
@@ -286,3 +285,61 @@ async fn test_get_commitment_results(mut mock_dependencies: MockDependencies) {
     assert_eq!(first_result.height, INITIAL_HEIGHT,);
     assert_eq!(second_result.height, INITIAL_HEIGHT.next().unwrap(),);
 }
+
+/// Adds two commitments and a revert task to the last commit and inserts the results into the
+/// channel. Returns the resulted height.
+async fn add_commitments_and_revert_tasks(
+    commitment_manager: &mut MockCommitmentManager,
+    mut height: BlockNumber,
+) -> BlockNumber {
+    for _ in 0..2 {
+        commitment_manager
+            .add_commitment_task(height, test_state_diff(), Some(StateDiffCommitment::default()))
+            .await
+            .unwrap();
+        height = height.next().unwrap();
+    }
+    height = height.prev().unwrap();
+    commitment_manager.add_revert_task(height, test_state_diff()).await.unwrap();
+
+    for _ in 0..3 {
+        commitment_manager.state_committer.pop_task_and_insert_result().await;
+    }
+    height
+}
+
+#[rstest]
+#[tokio::test]
+async fn test_wait_for_revert(mut mock_dependencies: MockDependencies) {
+    add_initial_heights(&mut mock_dependencies);
+    let commitment_manager_config = CommitmentManagerConfig {
+        tasks_channel_size: 5,
+        results_channel_size: 5,
+        wait_for_tasks_channel: false,
+    };
+    let mut commitment_manager =
+        create_mock_commitment_manager(mock_dependencies, Some(commitment_manager_config)).await;
+
+    let height = add_commitments_and_revert_tasks(&mut commitment_manager, INITIAL_HEIGHT).await;
+    let (commitment_results, revert_result) = commitment_manager.wait_for_revert_result().await;
+    assert_eq!(commitment_results.len(), 2);
+    assert_eq!(revert_result.height, height);
+}
+
+#[rstest]
+#[tokio::test]
+#[should_panic(expected = "Got revert output")]
+async fn test_revert_result_at_getting_commitments(mut mock_dependencies: MockDependencies) {
+    add_initial_heights(&mut mock_dependencies);
+    let commitment_manager_config = CommitmentManagerConfig {
+        tasks_channel_size: 5,
+        results_channel_size: 5,
+        wait_for_tasks_channel: false,
+    };
+    let mut commitment_manager =
+        create_mock_commitment_manager(mock_dependencies, Some(commitment_manager_config)).await;
+
+    add_commitments_and_revert_tasks(&mut commitment_manager, INITIAL_HEIGHT).await;
+    wait_for_filling_channel(&mut commitment_manager.results_receiver, 3).await;
+    commitment_manager.get_commitment_results().await;
+}

crates/apollo_batcher/src/test_utils.rs

@@ -5,8 +5,12 @@ use std::sync::Arc;
 use apollo_batcher_config::config::{BatcherConfig, FirstBlockWithPartialBlockHash};
 use apollo_batcher_types::batcher_types::{ProposalId, ProposeBlockInput};
 use apollo_class_manager_types::{EmptyClassManagerClient, SharedClassManagerClient};
-use apollo_committer_types::committer_types::CommitBlockResponse;
-use apollo_committer_types::communication::{MockCommitterClient, SharedCommitterClient};
+use apollo_committer_types::committer_types::{CommitBlockResponse, RevertBlockResponse};
+use apollo_committer_types::communication::{
+    CommitterRequest,
+    MockCommitterClient,
+    SharedCommitterClient,
+};
 use apollo_l1_provider_types::MockL1ProviderClient;
 use apollo_mempool_types::communication::MockMempoolClient;
 use apollo_mempool_types::mempool_types::CommitBlockArgs;
@@ -46,6 +50,7 @@ use crate::commitment_manager::types::{
     CommitmentTaskOutput,
     CommitterTaskInput,
     CommitterTaskOutput,
+    RevertTaskOutput,
 };
 use crate::pre_confirmed_block_writer::{
     MockPreconfirmedBlockWriterFactoryTrait,
@@ -345,10 +350,20 @@ impl MockStateCommitter {
     ) {
         while mock_task_receiver.recv().await.is_some() {
             let task = tasks_receiver.try_recv().unwrap();
-            let result = CommitterTaskOutput::Commit(CommitmentTaskOutput {
-                response: CommitBlockResponse { state_root: GlobalRoot::default() },
-                height: task.0.height(),
-            });
+            let result = match task.0 {
+                CommitterRequest::CommitBlock(commit_request) => {
+                    CommitterTaskOutput::Commit(CommitmentTaskOutput {
+                        response: CommitBlockResponse { state_root: GlobalRoot::default() },
+                        height: commit_request.height,
+                    })
+                }
+                CommitterRequest::RevertBlock(revert_request) => {
+                    CommitterTaskOutput::Revert(RevertTaskOutput {
+                        response: RevertBlockResponse::RevertedTo(GlobalRoot::default()),
+                        height: revert_request.height,
+                    })
+                }
+            };
             results_sender.try_send(result).unwrap();
         }
     }