From 459b9e7a4fbe4623bfa8d571b7ae16a14eea75b4 Mon Sep 17 00:00:00 2001
From: Anthony Alaribe <anthonyalaribe@gmail.com>
Date: Tue, 9 Jun 2026 19:42:17 +0200
Subject: [PATCH] test: fix test_flushable_buckets_carry_all_batches assertion
 drift
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The test asserted one RecordBatch per insert, but inserts cross
MAX_BATCH_COUNT_PER_BUCKET and trigger insert-time coalesce under the
bucket lock — so 10 single-row inserts come back as 2 RecordBatches,
not 10.

The flush path actually only cares that every row makes it through; the
batch-count breakdown is an implementation detail of the in-bucket
coalesce gate. Rewrite the assertion to check rows (no loss) plus a
bounded-batch-count sanity guard, and explain the coalesce trade-off
in the test comment so the next reader doesn't bring the loose
assertion back.

Was failing on master before this PR (verified via git stash); this
removes a CI gate that's been red on every PR.
---
 src/mem_buffer.rs | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/src/mem_buffer.rs b/src/mem_buffer.rs
index 5dbaf544..709aee24 100644
--- a/src/mem_buffer.rs
+++ b/src/mem_buffer.rs
@@ -2315,9 +2315,16 @@ mod tests {
 
     #[test]
     fn test_flushable_buckets_carry_all_batches() {
-        // We no longer pre-compact at flush time — the parquet writer downstream
-        // regroups rows into row groups itself, and pre-compacting forces an
-        // unnecessary deep copy of the entire bucket.
+        // Flush-time pre-compaction is gone (the parquet writer downstream
+        // regroups rows into row groups itself, and pre-compacting forced an
+        // unnecessary deep copy of the entire bucket).
+        //
+        // Insert-time coalesce still applies: when `batches.len()` crosses
+        // `MAX_BATCH_COUNT_PER_BUCKET` we concat under the bucket lock to
+        // amortise per-batch RecordBatch overhead. So after N tiny inserts
+        // the flushable bucket may carry fewer than N RecordBatches —
+        // the invariant the flush path actually cares about is "every row
+        // makes it through", not "one RecordBatch per insert".
         let buffer = MemBuffer::new();
         let ts = chrono::Utc::now().timestamp_micros();
 
@@ -2329,11 +2336,19 @@ mod tests {
 
         let cutoff = MemBuffer::compute_bucket_id(ts) + 1;
         let flushable = buffer.get_flushable_buckets(cutoff);
-        assert_eq!(flushable.len(), 1);
-        assert_eq!(flushable[0].batches.len(), total_rows);
+        assert_eq!(flushable.len(), 1, "all inserts share one time bucket");
         assert_eq!(flushable[0].row_count, total_rows);
         let summed: usize = flushable[0].batches.iter().map(|b| b.num_rows()).sum();
-        assert_eq!(summed, total_rows);
+        assert_eq!(summed, total_rows, "no rows lost to insert-time coalesce");
+        // Coalesce keeps the in-bucket batch count bounded by MAX_BATCH_COUNT_PER_BUCKET + 1
+        // (concat fires when len > MAX_BATCH_COUNT_PER_BUCKET, then the next push reaches
+        // the cap again before the next concat).
+        assert!(
+            flushable[0].batches.len() <= MAX_BATCH_COUNT_PER_BUCKET + 1,
+            "got {} batches, expected ≤ {}",
+            flushable[0].batches.len(),
+            MAX_BATCH_COUNT_PER_BUCKET + 1
+        );
     }
 
     #[test]