From d1deda5fdd00f0a9b03bda9a0c1be7e494c90089 Mon Sep 17 00:00:00 2001
From: Itamar Turner-Trauring <itamar@pythonspeed.com>
Date: Wed, 23 Nov 2022 10:09:53 -0500
Subject: [PATCH] Unit tests for _WriteBuffer.

---
 src/allmydata/test/test_storage.py | 47 +++++++++++++++++++++++++-----
 1 file changed, 39 insertions(+), 8 deletions(-)

diff --git a/src/allmydata/test/test_storage.py b/src/allmydata/test/test_storage.py
index 134609f81..f5762f616 100644
--- a/src/allmydata/test/test_storage.py
+++ b/src/allmydata/test/test_storage.py
@@ -3,14 +3,9 @@ Tests for allmydata.storage.
 
 Ported to Python 3.
 """
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from __future__ import unicode_literals
 
-from future.utils import native_str, PY2, bytes_to_native_str, bchr
-if PY2:
-    from future.builtins import filter, map, zip, ascii, chr, hex, input, next, oct, open, pow, round, super, bytes, dict, list, object, range, str, max, min  # noqa: F401
+from __future__ import annotations
+from future.utils import native_str, bytes_to_native_str, bchr
 from six import ensure_str
 
 from io import (
@@ -59,7 +54,7 @@ from allmydata.storage.common import storage_index_to_dir, \
      si_b2a, si_a2b
 from allmydata.storage.lease import LeaseInfo
 from allmydata.immutable.layout import WriteBucketProxy, WriteBucketProxy_v2, \
-     ReadBucketProxy
+     ReadBucketProxy, _WriteBuffer
 from allmydata.mutable.layout import MDMFSlotWriteProxy, MDMFSlotReadProxy, \
                                      LayoutInvalid, MDMFSIGNABLEHEADER, \
                                      SIGNED_PREFIX, MDMFHEADER, \
@@ -3746,3 +3741,39 @@ class LeaseInfoTests(SyncTestCase):
             info.to_mutable_data(),
             HasLength(info.mutable_size()),
         )
+
+
+class WriteBufferTests(SyncTestCase):
+    """Tests for ``_WriteBuffer``."""
+
+    @given(
+        small_writes=strategies.lists(
+            strategies.binary(min_size=1, max_size=20),
+            min_size=10, max_size=20),
+        batch_size=strategies.integers(min_value=5, max_value=10)
+    )
+    def test_write_buffer(self, small_writes: list[bytes], batch_size: int):
+        """
+        ``_WriteBuffer`` coalesces small writes into bigger writes based on
+        the batch size.
+        """
+        offset = 0
+        wb = _WriteBuffer(batch_size)
+        result = b""
+        for data in small_writes:
+            should_flush = wb.queue_write(offset, data)
+            offset += len(data)
+            if should_flush:
+                flushed_offset, flushed_data = wb.flush()
+                self.assertEqual(flushed_offset, len(result))
+                # The flushed data is in batch sizes, or closest approximation
+                # given queued inputs:
+                self.assertTrue(batch_size <= len(flushed_data) < batch_size + len(data))
+                result += flushed_data
+
+        # Final flush:
+        flushed_offset, flushed_data = wb.flush()
+        self.assertEqual(flushed_offset, len(result))
+        result += flushed_data
+
+        self.assertEqual(result, b"".join(small_writes))