From 1505223d820c946f401ae3ff695113419cb69e5a Mon Sep 17 00:00:00 2001
From: Dennis Lawler <4824647+drawlerr@users.noreply.github.com>
Date: Mon, 3 Feb 2020 09:49:09 -0800
Subject: [PATCH] Fix issue with fetching test-mode -1k document corpora.
 (#885)

---
 esrally/utils/net.py    | 18 ++++++++++++------
 tests/utils/net_test.py | 11 ++++++++++-
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/esrally/utils/net.py b/esrally/utils/net.py
index 5ae616853..a4a0190b8 100644
--- a/esrally/utils/net.py
+++ b/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):
diff --git a/tests/utils/net_test.py b/tests/utils/net_test.py
index 395d9d12e..c7aaecd58 100644
--- a/tests/utils/net_test.py
+++ b/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