Mercurial > public > sg101

annotate core/tests/test_download.py @ 1231:bbabbbb80d60 modernize tip

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Add unit test for S3 upload of images.
author Brian Neal <bgneal@gmail.com>
date Sun, 20 Apr 2025 19:15:39 -0500
parents c901261c6ce8
children
rev   line source
bgneal@1225 1 import unittest
bgneal@1225 2 from urlparse import urlparse, ParseResult
bgneal@1225 3
bgneal@1225 4 from mock import call, patch, mock_open, Mock
bgneal@1225 5
bgneal@1225 6 from core.download import download_file
bgneal@1225 7
bgneal@1225 8
bgneal@1225 9 class DownloadFileTestCase(unittest.TestCase):
bgneal@1225 10 """Unit tests for the download_file function."""
bgneal@1225 11
bgneal@1225 12 @patch('core.download.requests')
bgneal@1225 13 def test_get_throws_exception(self, requests_mock):
bgneal@1225 14 requests_mock.get.side_effect = RuntimeError('ope')
bgneal@1225 15 r = None
bgneal@1225 16 try:
bgneal@1225 17 r = download_file('url', path='path', timeout=5.0)
bgneal@1225 18 except RuntimeError:
bgneal@1225 19 pass
bgneal@1225 20 except Exception as e:
bgneal@1225 21 self.fail('Unexpected exception {}'.format(e))
bgneal@1225 22
bgneal@1225 23 self.assertIsNone(r)
bgneal@1225 24 self.assertEqual(requests_mock.get.mock_calls, [
bgneal@1225 25 call('url', stream=True, timeout=5.0),
bgneal@1225 26 ])
bgneal@1225 27
bgneal@1225 28 @patch('core.download.requests')
bgneal@1225 29 def test_get_returns_non_success(self, requests_mock):
bgneal@1225 30 response_mock = Mock()
bgneal@1225 31 response_mock.status_code = 404
bgneal@1225 32 requests_mock.get.return_value = response_mock
bgneal@1225 33 r = download_file('url', path='path', timeout=5.0)
bgneal@1225 34
bgneal@1225 35 self.assertIsNone(r)
bgneal@1225 36 self.assertEqual(requests_mock.get.mock_calls, [
bgneal@1225 37 call('url', stream=True, timeout=5.0),
bgneal@1225 38 ])
bgneal@1225 39
bgneal@1225 40 @patch('core.download.os.stat')
bgneal@1225 41 @patch('core.download.shutil.copyfileobj')
bgneal@1225 42 @patch('core.download.requests')
bgneal@1225 43 def test_happy_path_with_path_argument(self, requests_mock, copy_mock,
bgneal@1225 44 stat_mock):
bgneal@1225 45 response_mock = Mock()
bgneal@1225 46 response_mock.status_code = 200
bgneal@1225 47 requests_mock.get.return_value = response_mock
bgneal@1225 48
bgneal@1225 49 stat_mock.st_size = 512
bgneal@1225 50
bgneal@1225 51 open_mock = mock_open()
bgneal@1225 52 with patch('__builtin__.open', open_mock):
bgneal@1225 53 r = download_file('url', path='path', timeout=5.0)
bgneal@1225 54
bgneal@1225 55 self.assertEqual(r, 'path')
bgneal@1225 56 self.assertEqual(requests_mock.get.mock_calls, [
bgneal@1225 57 call('url', stream=True, timeout=5.0),
bgneal@1225 58 ])
bgneal@1225 59 self.assertTrue(response_mock.raw.decode_content)
bgneal@1225 60 self.assertEqual(copy_mock.mock_calls, [
bgneal@1225 61 call(response_mock.raw, open_mock.return_value),
bgneal@1225 62 ])
bgneal@1225 63
bgneal@1225 64 @patch('core.download.os.remove')
bgneal@1225 65 @patch('core.download.shutil.copyfileobj')
bgneal@1225 66 @patch('core.download.requests')
bgneal@1225 67 def test_copyfileobj_raises(self, requests_mock, copy_mock, remove_mock):
bgneal@1225 68 response_mock = Mock()
bgneal@1225 69 response_mock.status_code = 200
bgneal@1225 70 requests_mock.get.return_value = response_mock
bgneal@1225 71
bgneal@1225 72 copy_mock.side_effect = IOError
bgneal@1225 73
bgneal@1225 74 open_mock = mock_open()
bgneal@1225 75 with patch('__builtin__.open', open_mock):
bgneal@1225 76 try:
bgneal@1225 77 download_file('url', path='path', timeout=5.0)
bgneal@1225 78 except IOError:
bgneal@1225 79 pass
bgneal@1225 80 else:
bgneal@1225 81 self.fail('Should have thrown')
bgneal@1225 82
bgneal@1225 83 self.assertEqual(requests_mock.get.mock_calls, [
bgneal@1225 84 call('url', stream=True, timeout=5.0),
bgneal@1225 85 ])
bgneal@1225 86 self.assertTrue(response_mock.raw.decode_content)
bgneal@1225 87 self.assertEqual(copy_mock.mock_calls, [
bgneal@1225 88 call(response_mock.raw, open_mock.return_value),
bgneal@1225 89 ])
bgneal@1225 90 self.assertEqual(remove_mock.mock_calls, [
bgneal@1225 91 call('path'),
bgneal@1225 92 ])
bgneal@1225 93
bgneal@1225 94
bgneal@1225 95 @patch('core.download.os')
bgneal@1225 96 @patch('core.download.tempfile')
bgneal@1225 97 @patch('core.download.mimetypes')
bgneal@1225 98 @patch('core.download.os.remove')
bgneal@1225 99 @patch('core.download.shutil.copyfileobj')
bgneal@1225 100 @patch('core.download.requests')
bgneal@1225 101 def test_happy_path_with_suffix(self, requests_mock, copy_mock, remove_mock,
bgneal@1225 102 mime_mock, tempfile_mock, os_mock):
bgneal@1225 103 response_mock = Mock()
bgneal@1225 104 response_mock.status_code = 200
bgneal@1225 105 requests_mock.get.return_value = response_mock
bgneal@1225 106
bgneal@1225 107 response_mock.headers.get.return_value = 'image/jpeg'
bgneal@1225 108 mime_mock.guess_extension.return_value = '.jpe'
bgneal@1225 109 fd_mock = Mock()
bgneal@1225 110 tempfile_mock.mkstemp.return_value = (fd_mock, 'temp-path')
bgneal@1225 111
bgneal@1225 112 open_mock = mock_open()
bgneal@1225 113 with patch('__builtin__.open', open_mock):
bgneal@1225 114 download_file('url', timeout=5.0)
bgneal@1225 115
bgneal@1225 116 self.assertEqual(requests_mock.get.mock_calls, [
bgneal@1225 117 call('url', stream=True, timeout=5.0),
bgneal@1225 118 call().headers.get('content-type'),
bgneal@1225 119 ])
bgneal@1225 120 self.assertEqual(response_mock.headers.get.mock_calls, [
bgneal@1225 121 call('content-type'),
bgneal@1225 122 ])
bgneal@1225 123 self.assertEqual(mime_mock.guess_extension.mock_calls, [
bgneal@1225 124 call('image/jpeg'),
bgneal@1225 125 ])
bgneal@1225 126 self.assertEqual(tempfile_mock.mkstemp.mock_calls, [
bgneal@1225 127 call(suffix='.jpg'),
bgneal@1225 128 ])
bgneal@1225 129 self.assertEqual(os_mock.close.mock_calls, [
bgneal@1225 130 call(fd_mock),
bgneal@1225 131 ])
bgneal@1225 132 self.assertTrue(response_mock.raw.decode_content)
bgneal@1225 133 self.assertEqual(copy_mock.mock_calls, [
bgneal@1225 134 call(response_mock.raw, open_mock.return_value),
bgneal@1225 135 ])
bgneal@1225 136
bgneal@1225 137
bgneal@1225 138 @patch('core.download.urlparse')
bgneal@1225 139 @patch('core.download.os')
bgneal@1225 140 @patch('core.download.tempfile')
bgneal@1225 141 @patch('core.download.mimetypes')
bgneal@1225 142 @patch('core.download.os.remove')
bgneal@1225 143 @patch('core.download.shutil.copyfileobj')
bgneal@1225 144 @patch('core.download.requests')
bgneal@1225 145 def test_happy_path_with_no_suffix(self, requests_mock, copy_mock,
bgneal@1225 146 remove_mock, mime_mock, tempfile_mock,
bgneal@1225 147 os_mock, urlparse_mock):
bgneal@1225 148 response_mock = Mock()
bgneal@1225 149 response_mock.status_code = 200
bgneal@1225 150 requests_mock.get.return_value = response_mock
bgneal@1225 151
bgneal@1225 152 response_mock.headers.get.return_value = None
bgneal@1225 153 urlparse_mock.return_value = ParseResult(
bgneal@1225 154 scheme='https', netloc='www.example.com', path='/something.txt',
bgneal@1225 155 params='', query='', fragment='')
bgneal@1225 156 os_mock.path.splitext.return_value = ('/something', '.txt')
bgneal@1225 157 fd_mock = Mock()
bgneal@1225 158 tempfile_mock.mkstemp.return_value = (fd_mock, 'temp-path')
bgneal@1225 159
bgneal@1225 160 open_mock = mock_open()
bgneal@1225 161 with patch('__builtin__.open', open_mock):
bgneal@1225 162 download_file('url', timeout=5.0)
bgneal@1225 163
bgneal@1225 164 self.assertEqual(requests_mock.get.mock_calls, [
bgneal@1225 165 call('url', stream=True, timeout=5.0),
bgneal@1225 166 call().headers.get('content-type'),
bgneal@1225 167 ])
bgneal@1225 168 self.assertEqual(response_mock.headers.get.mock_calls, [
bgneal@1225 169 call('content-type'),
bgneal@1225 170 ])
bgneal@1225 171 self.assertEqual(urlparse_mock.mock_calls, [
bgneal@1225 172 call('url'),
bgneal@1225 173 ])
bgneal@1225 174 self.assertEqual(mime_mock.guess_extension.mock_calls, [])
bgneal@1225 175 self.assertEqual(os_mock.path.splitext.mock_calls, [
bgneal@1225 176 call('/something.txt'),
bgneal@1225 177 ])
bgneal@1225 178 self.assertEqual(tempfile_mock.mkstemp.mock_calls, [
bgneal@1225 179 call(suffix='.txt'),
bgneal@1225 180 ])
bgneal@1225 181 self.assertEqual(os_mock.close.mock_calls, [
bgneal@1225 182 call(fd_mock),
bgneal@1225 183 ])
bgneal@1225 184 self.assertTrue(response_mock.raw.decode_content)
bgneal@1225 185 self.assertEqual(copy_mock.mock_calls, [
bgneal@1225 186 call(response_mock.raw, open_mock.return_value),
bgneal@1225 187 ])