From a32616b50c456ace7423eb93632da765bdaa2495 Mon Sep 17 00:00:00 2001
From: Yu Ming Zhu <yzhu@redhat.com>
Date: Jun 14 2022 10:58:00 +0000
Subject: unittest


---

diff --git a/tests/test_cli/test_download_file.py b/tests/test_cli/test_download_file.py
index 1907324..aa5aecd 100644
--- a/tests/test_cli/test_download_file.py
+++ b/tests/test_cli/test_download_file.py
@@ -36,8 +36,9 @@ class TestDownloadFile(unittest.TestCase):
         self.stdout = mock.patch('sys.stdout', new_callable=six.StringIO).start()
         self.stderr = mock.patch('sys.stderr', new_callable=six.StringIO).start()
         self.request_with_retry = mock.patch('koji.request_with_retry').start()
+        self.head = mock.patch('requests.head').start()
         # will work when contextlib.closing will be removed in future
-        self.request_with_retry = self.request_with_retry.return_value.__enter__
+        # self.request_with_retry = self.request_with_retry.return_value.__enter__
 
     def tearDown(self):
         mock.patch.stopall()
@@ -58,6 +59,10 @@ class TestDownloadFile(unittest.TestCase):
     @mock_open()
     def test_handle_download_file(self, m_open):
         self.reset_mock()
+        headers = mock.MagicMock()
+        rsp_head = self.head.return_value
+        rsp_head.status_code = 200
+        rsp_head.headers = {'Content-Length': '5' }
         response = mock.MagicMock()
         self.request_with_retry.return_value = response
         response.headers.get.return_value = '5' # content-length
@@ -159,6 +164,7 @@ class TestDownloadFileError(unittest.TestCase):
 
     @requests_mock.Mocker()
     def test_handle_download_file_error_404(self, m):
+        m.head('http://url')
         m.get("http://url", text='Not Found\n', status_code=404)
         with self.assertRaises(requests.HTTPError):
             download_file("http://url", self.filename)
@@ -169,6 +175,7 @@ class TestDownloadFileError(unittest.TestCase):
 
     @requests_mock.Mocker()
     def test_handle_download_file_error_500(self, m):
+        m.head('http://url')
         m.get("http://url", text='Internal Server Error\n', status_code=500)
         with self.assertRaises(requests.HTTPError):
             download_file("http://url", self.filename)