From d30008884c9210c3c7032804219da03686d4c030 Mon Sep 17 00:00:00 2001
From: Balaji Arun <balaji@aptoslabs.com>
Date: Tue, 18 Jun 2024 19:51:12 -0700
Subject: [PATCH] [consensus][buffer_manager_test] compare all ordered blocks

---
 .../pipeline/tests/buffer_manager_tests.rs    | 25 ++++++++++++-------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/consensus/src/pipeline/tests/buffer_manager_tests.rs b/consensus/src/pipeline/tests/buffer_manager_tests.rs
index ed1247bc4c2c2..0f2fd3e5b1650 100644
--- a/consensus/src/pipeline/tests/buffer_manager_tests.rs
+++ b/consensus/src/pipeline/tests/buffer_manager_tests.rs
@@ -255,16 +255,23 @@ async fn assert_results(
     batches: Vec<Vec<PipelinedBlock>>,
     result_rx: &mut Receiver<OrderedBlocks>,
 ) {
-    for (i, batch) in enumerate(batches) {
+    let mut blocks: Vec<PipelinedBlock> = Vec::new();
+    for _ in 0..batches.len() {
         let OrderedBlocks { ordered_blocks, .. } = result_rx.next().await.unwrap();
-        assert_eq!(
-            ordered_blocks.last().unwrap().id(),
-            batch.last().unwrap().id(),
-            "Inconsistent Block IDs (expected {} got {}) for {}-th block",
-            batch.last().unwrap().id(),
-            ordered_blocks.last().unwrap().id(),
-            i,
-        );
+        blocks.extend(ordered_blocks.into_iter());
+    }
+
+    for (i, batch) in enumerate(batches) {
+        for (idx, ordered_block) in blocks.drain(..batch.len()).enumerate() {
+            assert_eq!(
+                ordered_block.id(),
+                batch[idx].id(),
+                "Inconsistent Block IDs (expected {} got {}) for {}-th block",
+                batch[idx].id(),
+                ordered_block.id(),
+                i,
+            );
+        }
     }
 }