From ba06d09f1408670a6c700481bdec17f0523f45a7 Mon Sep 17 00:00:00 2001
From: Balaji Arun <balaji@aptoslabs.com>
Date: Thu, 27 Jun 2024 16:01:05 -0700
Subject: [PATCH] [consensus][buffer_manager_test][bugfix] count batches
 properly (#13839)

---
 consensus/src/pipeline/tests/buffer_manager_tests.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/consensus/src/pipeline/tests/buffer_manager_tests.rs b/consensus/src/pipeline/tests/buffer_manager_tests.rs
index 919c96dcbccb6..8e6aee8bbed18 100644
--- a/consensus/src/pipeline/tests/buffer_manager_tests.rs
+++ b/consensus/src/pipeline/tests/buffer_manager_tests.rs
@@ -255,8 +255,9 @@ async fn assert_results(
     batches: Vec<Vec<PipelinedBlock>>,
     result_rx: &mut Receiver<OrderedBlocks>,
 ) {
+    let total_batches = batches.iter().flatten().count();
     let mut blocks: Vec<PipelinedBlock> = Vec::new();
-    for _ in 0..batches.len() {
+    while blocks.len() < total_batches {
         let OrderedBlocks { ordered_blocks, .. } = result_rx.next().await.unwrap();
         blocks.extend(ordered_blocks.into_iter());
     }
@@ -340,7 +341,6 @@ fn buffer_manager_happy_path_test() {
     });
 }
 
-#[ignore] // TODO: turn this test back on once the flakes have resolved.
 #[test]
 fn buffer_manager_sync_test() {
     // happy path