1# Copyright 2021 Google LLC 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14"""Tests for http_utils.py""" 15 16import unittest 17from unittest import mock 18 19from pyfakefs import fake_filesystem_unittest 20 21import http_utils 22 23mock_get_response = mock.MagicMock(status_code=200, content=b'') 24 25 26class DownloadUrlTest(unittest.TestCase): 27 """Tests that download_url works.""" 28 URL = 'https://example.com/file' 29 FILE_PATH = '/tmp/file' 30 31 @mock.patch('time.sleep') 32 @mock.patch('requests.get', return_value=mock_get_response) 33 def test_download_url_no_error(self, mock_urlretrieve, _): 34 """Tests that download_url works when there is no error.""" 35 self.assertTrue(http_utils.download_url(self.URL, self.FILE_PATH)) 36 self.assertEqual(1, mock_urlretrieve.call_count) 37 38 @mock.patch('time.sleep') 39 @mock.patch('logging.error') 40 @mock.patch('requests.get', 41 return_value=mock.MagicMock(status_code=404, content=b'')) 42 def test_download_url_http_error(self, mock_get, mock_error, _): 43 """Tests that download_url doesn't retry when there is an HTTP error.""" 44 self.assertFalse(http_utils.download_url(self.URL, self.FILE_PATH)) 45 mock_error.assert_called_with( 46 'Unable to download from: %s. Code: %d. Content: %s.', self.URL, 404, 47 b'') 48 self.assertEqual(1, mock_get.call_count) 49 50 @mock.patch('time.sleep') 51 @mock.patch('requests.get', side_effect=ConnectionResetError) 52 def test_download_url_connection_error(self, mock_get, mock_sleep): 53 """Tests that download_url doesn't retry when there is an HTTP error.""" 54 self.assertFalse(http_utils.download_url(self.URL, self.FILE_PATH)) 55 self.assertEqual(4, mock_get.call_count) 56 self.assertEqual(3, mock_sleep.call_count) 57 58 59class DownloadAndUnpackZipTest(fake_filesystem_unittest.TestCase): 60 """Tests download_and_unpack_zip.""" 61 62 def setUp(self): 63 self.setUpPyfakefs() 64 65 @mock.patch('requests.get', return_value=mock_get_response) 66 def test_bad_zip_download(self, _): 67 """Tests download_and_unpack_zip returns none when a bad zip is passed.""" 68 self.fs.create_file('/url_tmp.zip', contents='Test file.') 69 self.assertFalse( 70 http_utils.download_and_unpack_zip('/not/a/real/url', 71 '/extract-directory')) 72