From 154b8c6021377d147e64ce489a63db31900d707f Mon Sep 17 00:00:00 2001
From: Jana Cupova <root@localhost>
Date: Jun 13 2022 11:29:22 +0000
Subject: koji download-task retry download file


Fixes: https://pagure.io/koji/issue/3375

---

diff --git a/cli/koji_cli/lib.py b/cli/koji_cli/lib.py
index 04ef329..c9c7ba5 100644
--- a/cli/koji_cli/lib.py
+++ b/cli/koji_cli/lib.py
@@ -568,6 +568,11 @@ def download_file(url, relpath, quiet=False, noprogress=False, size=None,
         else:
             print("Downloading: %s" % relpath)
 
+    if not filesize:
+        response = requests.head(url, timeout=10)
+        if response.status_code == 200 and response.headers.get('Content-Length'):
+            filesize = int(response.headers['Content-Length'])
+
     pos = 0
     headers = {}
     if filesize:
@@ -588,7 +593,7 @@ def download_file(url, relpath, quiet=False, noprogress=False, size=None,
 
     try:
         # closing needs to be used for requests < 2.18.0
-        with closing(requests.get(url, headers=headers, stream=True)) as response:
+        with closing(koji.request_with_retry().get(url, headers=headers, stream=True)) as response:
             if response.status_code in (200, 416):  # full content provided or reaching behind EOF
                 # rewrite in such case
                 f.close()
diff --git a/tests/test_cli/test_download_file.py b/tests/test_cli/test_download_file.py
index 894398f..1907324 100644
--- a/tests/test_cli/test_download_file.py
+++ b/tests/test_cli/test_download_file.py
@@ -7,6 +7,7 @@ import os
 import requests_mock
 import requests
 import unittest
+import koji
 
 from koji_cli.lib import download_file, _download_progress
 
@@ -27,16 +28,16 @@ class TestDownloadFile(unittest.TestCase):
         self.stdout.truncate()
         self.stderr.seek(0)
         self.stderr.truncate()
-        self.requests_get.reset_mock()
+        self.request_with_retry.reset_mock()
 
     def setUp(self):
         self.tempdir = tempfile.mkdtemp()
         self.filename = self.tempdir + "/filename"
         self.stdout = mock.patch('sys.stdout', new_callable=six.StringIO).start()
         self.stderr = mock.patch('sys.stderr', new_callable=six.StringIO).start()
-        self.requests_get = mock.patch('requests.get', create=True, name='requests.get').start()
+        self.request_with_retry = mock.patch('koji.request_with_retry').start()
         # will work when contextlib.closing will be removed in future
-        #self.requests_get = self.requests_get.return_value.__enter__
+        self.request_with_retry = self.request_with_retry.return_value.__enter__
 
     def tearDown(self):
         mock.patch.stopall()
@@ -58,7 +59,7 @@ class TestDownloadFile(unittest.TestCase):
     def test_handle_download_file(self, m_open):
         self.reset_mock()
         response = mock.MagicMock()
-        self.requests_get.return_value = response
+        self.request_with_retry.return_value = response
         response.headers.get.return_value = '5' # content-length
         response.iter_content.return_value = ['abcde']
 
@@ -68,7 +69,7 @@ class TestDownloadFile(unittest.TestCase):
         expected = 'Downloading: %s\n[====================================] 100%%     5.00 B\r\n' % self.filename
         self.assertMultiLineEqual(actual, expected)
 
-        self.requests_get.assert_called_once()
+        self.request_with_retry.assert_called_once()
         m_open.assert_called_once()
         response.headers.get.assert_called_once()
         response.iter_content.assert_called_once()
@@ -78,7 +79,7 @@ class TestDownloadFile(unittest.TestCase):
     def test_handle_download_file_undefined_length(self, m_open):
         self.reset_mock()
         response = mock.MagicMock()
-        self.requests_get.return_value = response
+        self.request_with_retry.return_value = response
         response.headers.get.return_value = None # content-length
         response.iter_content.return_value = ['a' * 65536, 'b' * 65536]
 
@@ -88,7 +89,7 @@ class TestDownloadFile(unittest.TestCase):
         expected = 'Downloading: %s\n[                                    ] ???%%  64.00 KiB\r[                                    ] ???%% 128.00 KiB\r[====================================] 100%% 128.00 KiB\r\n' % self.filename
         self.assertMultiLineEqual(actual, expected)
 
-        self.requests_get.assert_called_once()
+        self.request_with_retry.assert_called_once()
         m_open.assert_called_once()
         response.headers.get.assert_called_once()
         response.iter_content.assert_called_once()
@@ -100,7 +101,7 @@ class TestDownloadFile(unittest.TestCase):
         actual = self.stdout.getvalue()
         expected = 'Downloading [8/10]: %s\n\n' % self.filename
         self.assertMultiLineEqual(actual, expected)
-        self.requests_get.assert_called_once()
+        self.request_with_retry.assert_called_once()
         self.assertIsNone(rv)
 
     def test_handle_download_file_quiet_noprogress(self):