Mercurial > public > sg101

comparison core/tests/test_download.py @ 1225:c901261c6ce8 modernize

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