use moto for tests

Author: khvn26

api/import_export/utils.py

@@ -61,8 +61,9 @@ def __exit__(
                 logger.warning("Aborted multipart upload due to error: %s", exc_val)
             return
 
-        # Upload any remaining data in the buffer
-        if self._buffer.tell() > 0:
+        # Upload any remaining data in the buffer (or an empty part if no data)
+        # S3 requires at least one part to complete a multipart upload
+        if self._buffer.tell() > 0 or not self._parts:
             self._upload_part()
 
         assert self._upload_id

api/tests/unit/import_export/test_unit_import_export_utils.py

@@ -1,53 +1,36 @@
-from unittest.mock import MagicMock, call
-
-import pytest
+import boto3
+from moto import mock_s3  # type: ignore[import-untyped]
+from pytest_mock import MockerFixture
 
 from import_export.utils import S3MultipartUploadWriter
 
 
-@pytest.fixture
-def s3_client() -> MagicMock:
-    client = MagicMock()
-    client.create_multipart_upload.return_value = {"UploadId": "test-upload-id"}
-    client.upload_part.return_value = {"ETag": "test-etag"}
-    return client
-
-
-def test_s3_multipart_upload_writer__single_part__completes_upload(
-    s3_client: MagicMock,
-) -> None:
+@mock_s3  # type: ignore[misc]
+def test_s3_multipart_upload_writer__single_part__completes_upload() -> None:
     # Given
     bucket_name = "test-bucket"
     key = "test-key"
     data = b"small data"
 
+    s3_resource = boto3.resource("s3", region_name="eu-west-2")
+    s3_resource.create_bucket(
+        Bucket=bucket_name,
+        CreateBucketConfiguration={"LocationConstraint": "eu-west-2"},
+    )
+    s3_client = boto3.client("s3", region_name="eu-west-2")
+
     # When
     with S3MultipartUploadWriter(s3_client, bucket_name, key) as writer:
         writer.write(data)
 
     # Then
-    s3_client.create_multipart_upload.assert_called_once_with(
-        Bucket=bucket_name,
-        Key=key,
-    )
-    s3_client.upload_part.assert_called_once_with(
-        Bucket=bucket_name,
-        Key=key,
-        PartNumber=1,
-        UploadId="test-upload-id",
-        Body=data,
-    )
-    s3_client.complete_multipart_upload.assert_called_once_with(
-        Bucket=bucket_name,
-        Key=key,
-        UploadId="test-upload-id",
-        MultipartUpload={"Parts": [{"PartNumber": 1, "ETag": "test-etag"}]},
-    )
-    s3_client.abort_multipart_upload.assert_not_called()
+    result = s3_client.get_object(Bucket=bucket_name, Key=key)
+    assert result["Body"].read() == data
 
 
+@mock_s3  # type: ignore[misc]
 def test_s3_multipart_upload_writer__multiple_parts__uploads_each_part(
-    s3_client: MagicMock,
+    mocker: MockerFixture,
 ) -> None:
     # Given
     bucket_name = "test-bucket"
@@ -58,11 +41,13 @@ def test_s3_multipart_upload_writer__multiple_parts__uploads_each_part(
     second_chunk = b"b" * chunk_size
     final_chunk = b"final"
 
-    s3_client.upload_part.side_effect = [
-        {"ETag": "etag-1"},
-        {"ETag": "etag-2"},
-        {"ETag": "etag-3"},
-    ]
+    s3_resource = boto3.resource("s3", region_name="eu-west-2")
+    s3_resource.create_bucket(
+        Bucket=bucket_name,
+        CreateBucketConfiguration={"LocationConstraint": "eu-west-2"},
+    )
+    s3_client = boto3.client("s3", region_name="eu-west-2")
+    upload_part_spy = mocker.spy(s3_client, "upload_part")
 
     # When
     with S3MultipartUploadWriter(s3_client, bucket_name, key) as writer:
@@ -71,134 +56,123 @@ def test_s3_multipart_upload_writer__multiple_parts__uploads_each_part(
         writer.write(final_chunk)
 
     # Then
-    assert s3_client.upload_part.call_count == 3
-    s3_client.upload_part.assert_has_calls(
-        [
-            call(
-                Bucket=bucket_name,
-                Key=key,
-                PartNumber=1,
-                UploadId="test-upload-id",
-                Body=first_chunk,
-            ),
-            call(
-                Bucket=bucket_name,
-                Key=key,
-                PartNumber=2,
-                UploadId="test-upload-id",
-                Body=second_chunk,
-            ),
-            call(
-                Bucket=bucket_name,
-                Key=key,
-                PartNumber=3,
-                UploadId="test-upload-id",
-                Body=final_chunk,
-            ),
-        ]
-    )
-    s3_client.complete_multipart_upload.assert_called_once_with(
-        Bucket=bucket_name,
-        Key=key,
-        UploadId="test-upload-id",
-        MultipartUpload={
-            "Parts": [
-                {"PartNumber": 1, "ETag": "etag-1"},
-                {"PartNumber": 2, "ETag": "etag-2"},
-                {"PartNumber": 3, "ETag": "etag-3"},
-            ]
-        },
-    )
+    result = s3_client.get_object(Bucket=bucket_name, Key=key)
+    assert result["Body"].read() == first_chunk + second_chunk + final_chunk
+
+    # Verify exactly 3 parts were uploaded
+    assert upload_part_spy.call_count == 3
+    # Verify part numbers are sequential
+    part_numbers = [
+        call.kwargs["PartNumber"] for call in upload_part_spy.call_args_list
+    ]
+    assert part_numbers == [1, 2, 3]
 
 
-def test_s3_multipart_upload_writer__accumulates_small_writes__uploads_when_threshold_reached(
-    s3_client: MagicMock,
+@mock_s3  # type: ignore[misc]
+def test_s3_multipart_upload_writer__accumulates_small_writes__uploads_correctly(
+    mocker: MockerFixture,
 ) -> None:
     # Given
     bucket_name = "test-bucket"
     key = "test-key"
     small_chunk = b"x" * 1000  # 1KB
     writes_to_reach_threshold = (S3MultipartUploadWriter.MIN_PART_SIZE // 1000) + 1
 
-    s3_client.upload_part.side_effect = [
-        {"ETag": "etag-1"},
-        {"ETag": "etag-2"},
-    ]
+    s3_resource = boto3.resource("s3", region_name="eu-west-2")
+    s3_resource.create_bucket(
+        Bucket=bucket_name,
+        CreateBucketConfiguration={"LocationConstraint": "eu-west-2"},
+    )
+    s3_client = boto3.client("s3", region_name="eu-west-2")
+    upload_part_spy = mocker.spy(s3_client, "upload_part")
 
     # When
     with S3MultipartUploadWriter(s3_client, bucket_name, key) as writer:
         for _ in range(writes_to_reach_threshold):
             writer.write(small_chunk)
-        # Write one more small chunk to have remaining data
         writer.write(b"final")
 
     # Then
-    # Should have uploaded one part when threshold was reached,
-    # and one final part with remaining data on exit
-    assert s3_client.upload_part.call_count == 2
+    result = s3_client.get_object(Bucket=bucket_name, Key=key)
+    expected_data = (small_chunk * writes_to_reach_threshold) + b"final"
+    assert result["Body"].read() == expected_data
 
+    # Verify buffering: one part when threshold reached, one final part on exit
+    assert upload_part_spy.call_count == 2
 
-def test_s3_multipart_upload_writer__error_during_write__aborts_upload(
-    s3_client: MagicMock,
-) -> None:
+
+@mock_s3  # type: ignore[misc]
+def test_s3_multipart_upload_writer__error_during_write__aborts_upload() -> None:
     # Given
     bucket_name = "test-bucket"
     key = "test-key"
 
+    s3_resource = boto3.resource("s3", region_name="eu-west-2")
+    s3_resource.create_bucket(
+        Bucket=bucket_name,
+        CreateBucketConfiguration={"LocationConstraint": "eu-west-2"},
+    )
+    s3_client = boto3.client("s3", region_name="eu-west-2")
+
     # When
-    with pytest.raises(ValueError, match="test error"):
+    try:
         with S3MultipartUploadWriter(s3_client, bucket_name, key) as writer:
             writer.write(b"some data")
             raise ValueError("test error")
+    except ValueError:
+        pass
 
-    # Then
-    s3_client.abort_multipart_upload.assert_called_once_with(
-        Bucket=bucket_name,
-        Key=key,
-        UploadId="test-upload-id",
-    )
-    s3_client.complete_multipart_upload.assert_not_called()
+    # Then - the object should not exist (upload was aborted)
+    objects = s3_client.list_objects_v2(Bucket=bucket_name)
+    assert objects.get("KeyCount", 0) == 0
 
 
-def test_s3_multipart_upload_writer__no_data__completes_with_no_parts(
-    s3_client: MagicMock,
-) -> None:
+@mock_s3  # type: ignore[misc]
+def test_s3_multipart_upload_writer__no_data__completes_with_empty_object() -> None:
     # Given
     bucket_name = "test-bucket"
     key = "test-key"
 
+    s3_resource = boto3.resource("s3", region_name="eu-west-2")
+    s3_resource.create_bucket(
+        Bucket=bucket_name,
+        CreateBucketConfiguration={"LocationConstraint": "eu-west-2"},
+    )
+    s3_client = boto3.client("s3", region_name="eu-west-2")
+
     # When
     with S3MultipartUploadWriter(s3_client, bucket_name, key):
         pass  # No writes
 
     # Then
-    s3_client.upload_part.assert_not_called()
-    s3_client.complete_multipart_upload.assert_called_once_with(
-        Bucket=bucket_name,
-        Key=key,
-        UploadId="test-upload-id",
-        MultipartUpload={"Parts": []},
-    )
+    result = s3_client.get_object(Bucket=bucket_name, Key=key)
+    assert result["Body"].read() == b""
 
 
+@mock_s3  # type: ignore[misc]
 def test_s3_multipart_upload_writer__multiple_small_writes__buffers_correctly(
-    s3_client: MagicMock,
+    mocker: MockerFixture,
 ) -> None:
     # Given
     bucket_name = "test-bucket"
     key = "test-key"
 
+    s3_resource = boto3.resource("s3", region_name="eu-west-2")
+    s3_resource.create_bucket(
+        Bucket=bucket_name,
+        CreateBucketConfiguration={"LocationConstraint": "eu-west-2"},
+    )
+    s3_client = boto3.client("s3", region_name="eu-west-2")
+    upload_part_spy = mocker.spy(s3_client, "upload_part")
+
     # When
     with S3MultipartUploadWriter(s3_client, bucket_name, key) as writer:
         writer.write(b"hello ")
         writer.write(b"world")
 
     # Then
-    # Both writes should be buffered and uploaded as single part
-    s3_client.upload_part.assert_called_once_with(
-        Bucket=bucket_name,
-        Key=key,
-        PartNumber=1,
-        UploadId="test-upload-id",
-        Body=b"hello world",
-    )
+    result = s3_client.get_object(Bucket=bucket_name, Key=key)
+    assert result["Body"].read() == b"hello world"
+
+    # Verify both writes were buffered and uploaded as a single part
+    assert upload_part_spy.call_count == 1