Merge pull request #2168 from Sanjays2402/fix/s3-stream-chunk-size

[S3] Respect chunk_size in download_object_as_stream

CHANGES.rst

diff --git a/CHANGES.rst b/CHANGES.rst
index 5b208a8..f1538e8 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -57,11 +57,20 @@
 Storage
 ~~~~~~~
 
+- [S3] Fix ``chunk_size`` argument being ignored by
+  ``download_object_as_stream`` and ``download_object_range_as_stream``. The
+  requested chunk size is now passed through to the underlying iterator
+  instead of the hardcoded default.
+  (GITHUB-1798)
+
+  [Sanjay Santhanam - @Sanjays2402]
+
 - [Azure Blobs] Fix ``chunk_size`` argument being ignored by
   ``download_object_as_stream`` and ``download_object_range_as_stream``. The
   requested chunk size is now forwarded to the underlying iterator instead of
   always using ``AZURE_DOWNLOAD_CHUNK_SIZE``.
   (GITHUB-1698)
+
   [Sanjay Santhanam - @Sanjays2402]
 
 DNS

libcloud/storage/drivers/s3.py

diff --git a/libcloud/storage/drivers/s3.py b/libcloud/storage/drivers/s3.py
index bd9212e..4b59e2d 100644
--- a/libcloud/storage/drivers/s3.py
+++ b/libcloud/storage/drivers/s3.py
@@ -521,7 +521,7 @@
             callback=read_in_chunks,
             response=response,
             callback_kwargs={
-                "iterator": response.iter_content(CHUNK_SIZE),
+                "iterator": response.iter_content(chunk_size or CHUNK_SIZE),
                 "chunk_size": chunk_size,
             },
             success_status_code=httplib.OK,
@@ -573,7 +573,7 @@
             callback=read_in_chunks,
             response=response,
             callback_kwargs={
-                "iterator": response.iter_content(CHUNK_SIZE),
+                "iterator": response.iter_content(chunk_size or CHUNK_SIZE),
                 "chunk_size": chunk_size,
             },
             success_status_code=httplib.PARTIAL_CONTENT,

libcloud/test/storage/test_s3.py

diff --git a/libcloud/test/storage/test_s3.py b/libcloud/test/storage/test_s3.py
index be70898..f8b4607 100644
--- a/libcloud/test/storage/test_s3.py
+++ b/libcloud/test/storage/test_s3.py
@@ -861,6 +861,59 @@
         finally:
             self.driver_type._get_object = old_func
 
+    def test_download_object_as_stream_uses_chunk_size(self):
+        # Regression test: the chunk_size passed by the caller must be
+        # forwarded to iter_content instead of the hardcoded CHUNK_SIZE.
+        container = Container(name="foo_bar_container", extra={}, driver=self.driver)
+        obj = Object(
+            name="foo_bar_object",
+            size=1000,
+            hash=None,
+            extra={},
+            container=container,
+            meta_data=None,
+            driver=self.driver_type,
+        )
+
+        requested_chunk_size = CHUNK_SIZE * 2
+        mock_response = Mock(name="mock response")
+        mock_response.iter_content.return_value = iter([b"a"])
+
+        with mock.patch.object(self.driver.connection, "request", return_value=mock_response):
+            with mock.patch.object(self.driver, "_get_object", side_effect=lambda **kw: kw):
+                self.driver.download_object_as_stream(obj=obj, chunk_size=requested_chunk_size)
+
+        mock_response.iter_content.assert_called_once_with(requested_chunk_size)
+
+    def test_download_object_range_as_stream_uses_chunk_size(self):
+        # Same regression as above, but for the ranged variant which goes
+        # through the same _get_object() code path.
+        container = Container(name="foo_bar_container", extra={}, driver=self.driver)
+        obj = Object(
+            name="foo_bar_object",
+            size=1000,
+            hash=None,
+            extra={},
+            container=container,
+            meta_data=None,
+            driver=self.driver_type,
+        )
+
+        requested_chunk_size = CHUNK_SIZE * 2
+        mock_response = Mock(name="mock response")
+        mock_response.iter_content.return_value = iter([b"a"])
+
+        with mock.patch.object(self.driver.connection, "request", return_value=mock_response):
+            with mock.patch.object(self.driver, "_get_object", side_effect=lambda **kw: kw):
+                self.driver.download_object_range_as_stream(
+                    obj=obj,
+                    start_bytes=0,
+                    end_bytes=100,
+                    chunk_size=requested_chunk_size,
+                )
+
+        mock_response.iter_content.assert_called_once_with(requested_chunk_size)
+
     def test_upload_object_invalid_ex_storage_class(self):
         # Invalid hash is detected on the amazon side and BAD_REQUEST is
         # returned