From 6eedae1bd5bd36f4c8e8e0648ca5cb331bbead22 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Patrick=20St=C3=A4hlin?= <patrick.staehlin@aiven.io>
Date: Wed, 20 Mar 2024 22:17:22 +0100
Subject: [PATCH] Fix multipart S3 uploads

Not sure how this worked before but something broke multipart uploads.
This adds a test to ensure that it keeps working and also tries to fix
the issue.
---
 rohmu/object_storage/s3.py     |  1 -
 test/object_storage/test_s3.py | 23 +++++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/rohmu/object_storage/s3.py b/rohmu/object_storage/s3.py
index a735fc52..cbc4a858 100644
--- a/rohmu/object_storage/s3.py
+++ b/rohmu/object_storage/s3.py
@@ -447,7 +447,6 @@ def multipart_upload_file_object(
                 if progress_fn:
                     # TODO: change this to incremental progress. Size parameter is currently unused.
                     progress_fn(bytes_sent, size)  # type: ignore[arg-type]
-                break
 
         self.stats.operation(StorageOperation.multipart_complete)
         try:
diff --git a/test/object_storage/test_s3.py b/test/object_storage/test_s3.py
index 55475a28..da21afcc 100644
--- a/test/object_storage/test_s3.py
+++ b/test/object_storage/test_s3.py
@@ -64,6 +64,29 @@ def test_store_file_from_disk(infra: S3Infra) -> None:
     )
 
 
+def test_store_file_object_large(infra: S3Infra) -> None:
+    test_data = b"test-data" * 2
+    chunk_size = len(test_data) // 2
+    file_object = BytesIO(test_data)
+
+    infra.transfer.multipart_chunk_size = chunk_size  # simulate smaller chunk size to force multiple chunks
+
+    metadata = {"Content-Length": len(test_data), "some-date": datetime(2022, 11, 15, 18, 30, 58, 486644)}
+    infra.transfer.store_file_object(key="test_key2", fd=file_object, metadata=metadata, multipart=True)
+
+    notifier = infra.notifier
+    s3_client = infra.s3_client
+
+    s3_client.create_multipart_upload.assert_called()
+    assert s3_client.upload_part.call_count == 2
+    s3_client.complete_multipart_upload.assert_called()
+    notifier.object_created.assert_called_once_with(
+        key="test_key2",
+        size=len(test_data),
+        metadata={"Content-Length": "18", "some-date": "2022-11-15 18:30:58.486644"},
+    )
+
+
 @pytest.mark.parametrize("multipart", [False, None, True])
 def test_store_file_object(infra: S3Infra, multipart: Optional[bool]) -> None:
     test_data = b"test-data"