Fix issue with fetching test-mode -1k document corpora. (#885)

esrally/utils/net.py

@@ -58,9 +58,12 @@ def __init__(self, msg, accuracy=0):
 
     def __call__(self, bytes_read, bytes_total):
         from esrally.utils import convert
-        completed = bytes_read / bytes_total
-        total_as_mb = convert.bytes_to_human_string(bytes_total)
-        self.p.print("%s (%s total size)" % (self.msg, total_as_mb), self.percent_format % (completed * 100))
+        if bytes_total:
+            completed = bytes_read / bytes_total
+            total_as_mb = convert.bytes_to_human_string(bytes_total)
+            self.p.print("%s (%s total size)" % (self.msg, total_as_mb), self.percent_format % (completed * 100))
+        else:
+            self.p.print(self.msg, ".")
 
     def finish(self):
         self.p.finish()
@@ -82,10 +85,13 @@ def __call__(self, bytes_amount):
             self._progress(self._bytes_read, self._expected_size_in_bytes)
 
     s3 = boto3.resource("s3")
+    bucket = s3.Bucket(bucket_name)
+    if expected_size_in_bytes is None:
+        expected_size_in_bytes = bucket.Object(bucket_path).content_length
     progress_callback = S3ProgressAdapter(expected_size_in_bytes, progress_indicator) if progress_indicator else None
-    s3.Bucket(bucket_name).download_file(bucket_path, local_path,
-                                         Callback=progress_callback,
-                                         Config=boto3.s3.transfer.TransferConfig(use_threads=False))
+    bucket.download_file(bucket_path, local_path,
+                         Callback=progress_callback,
+                         Config=boto3.s3.transfer.TransferConfig(use_threads=False))
 
 
 def download_s3(url, local_path, expected_size_in_bytes=None, progress_indicator=None):

tests/utils/net_test.py

@@ -14,7 +14,6 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
 import random
 import unittest.mock as mock
 from unittest import TestCase
@@ -33,3 +32,13 @@ def test_download_from_s3_bucket(self, download):
                         expected_size, progress_indicator)
         download.assert_called_once_with("mybucket.elasticsearch.org", "data/documents.json.bz2",
                                          "/tmp/documents.json.bz2", expected_size, progress_indicator)
+
+    def test_progress(self):
+        progress = net.Progress("test")
+        mock_progress = mock.Mock()
+        progress.p = mock_progress
+        progress(42, 100)
+        assert mock_progress.print.called
+        mock_progress.reset_mock()
+        progress(42, None)
+        assert mock_progress.print.called