1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
|
try:
import unittest2 as unittest # python < 2.7 compat
except ImportError:
import unittest
import tempfile
import shutil
import os
import time
import six
from mock import Mock
from dlmanager import manager as download_manager
def mock_session():
response = Mock()
session = Mock(get=Mock(return_value=response))
return session, response
def mock_response(response, data, wait=0):
data = six.b(data)
def iter_content(chunk_size=4):
rest = data
while rest:
time.sleep(wait)
chunk = rest[:chunk_size]
rest = rest[chunk_size:]
yield chunk
response.headers = {'Content-length': str(len(data))}
response.iter_content = iter_content
class TestDownload(unittest.TestCase):
def setUp(self):
self.tempdir = tempfile.mkdtemp()
self.addCleanup(shutil.rmtree, self.tempdir)
self.finished = Mock()
self.session, self.session_response = mock_session()
self.tempfile = os.path.join(self.tempdir, 'dest')
self.dl = download_manager.Download('http://url', self.tempfile,
finished_callback=self.finished,
chunk_size=4,
session=self.session)
def test_creation(self):
self.assertFalse(self.dl.is_canceled())
self.assertFalse(self.dl.is_running())
self.assertIsNone(self.dl.error())
self.assertEquals(self.dl.get_url(), 'http://url')
self.assertEquals(self.dl.get_dest(), self.tempfile)
def create_response(self, data, wait=0):
mock_response(self.session_response, data, wait)
def test_download(self):
self.create_response('1234' * 4, 0.01)
# no file present yet
self.assertFalse(os.path.exists(self.tempfile))
self.dl.start()
self.assertTrue(self.dl.is_running())
self.dl.wait()
self.assertFalse(self.dl.is_running())
self.finished.assert_called_with(self.dl)
# file has been downloaded
with open(self.tempfile) as f:
self.assertEquals(f.read(), '1234' * 4)
def test_download_cancel(self):
self.create_response('1234' * 1000, wait=0.01)
start = time.time()
self.dl.start()
time.sleep(0.1)
self.dl.cancel()
with self.assertRaises(download_manager.DownloadInterrupt):
self.dl.wait()
self.assertTrue(self.dl.is_canceled())
# response generation should have taken 1000 * 0.01 = 10 seconds.
# since we canceled, this must be lower.
self.assertTrue((time.time() - start) < 1.0)
# file was deleted
self.assertFalse(os.path.exists(self.tempfile))
# finished callback was called
self.finished.assert_called_with(self.dl)
def test_download_with_progress(self):
data = []
def update_progress(_dl, current, total):
data.append((_dl, current, total))
self.create_response('1234' * 4)
self.dl.set_progress(update_progress)
self.dl.start()
self.dl.wait()
self.assertEquals(data, [
(self.dl, 0, 16),
(self.dl, 4, 16),
(self.dl, 8, 16),
(self.dl, 12, 16),
(self.dl, 16, 16),
])
# file has been downloaded
with open(self.tempfile) as f:
self.assertEquals(f.read(), '1234' * 4)
# finished callback was called
self.finished.assert_called_with(self.dl)
def test_download_error_in_thread(self):
self.session_response.headers = {'Content-length': '24'}
self.session_response.iter_content.side_effect = IOError
self.dl.start()
with self.assertRaises(IOError):
self.dl.wait()
self.assertEquals(self.dl.error()[0], IOError)
# finished callback was called
self.finished.assert_called_with(self.dl)
def test_wait_does_not_block_on_exception(self):
# this test the case when a user may hit CTRL-C for example
# during a dl.wait() call.
self.create_response('1234' * 1000, wait=0.01)
original_join = self.dl.thread.join
it = iter('123')
def join(timeout=None):
next(it) # will throw StopIteration after a few calls
original_join(timeout)
self.dl.thread.join = join
start = time.time()
self.dl.start()
with self.assertRaises(StopIteration):
self.dl.wait()
self.assertTrue(self.dl.is_canceled())
# wait for the thread to finish
original_join()
# response generation should have taken 1000 * 0.01 = 10 seconds.
# since we got an error, this must be lower.
self.assertTrue((time.time() - start) < 1.0)
# file was deleted
self.assertFalse(os.path.exists(self.tempfile))
# finished callback was called
self.finished.assert_called_with(self.dl)
class TestDownloadManager(unittest.TestCase):
def setUp(self):
self.tempdir = tempfile.mkdtemp()
self.addCleanup(shutil.rmtree, self.tempdir)
self.dl_manager = download_manager.DownloadManager(self.tempdir)
def do_download(self, url, fname, data, wait=0):
session, response = mock_session()
mock_response(response, data, wait)
# patch the session, so the download will use that
self.dl_manager.session = session
return self.dl_manager.download(url, fname)
def test_download(self):
dl1 = self.do_download('http://foo', 'foo', 'hello' * 4, wait=0.02)
self.assertIsInstance(dl1, download_manager.Download)
self.assertTrue(dl1.is_running())
# with the same fname, no new download is started. The same instance
# is returned since the download is running.
dl2 = self.do_download('http://bar', 'foo', 'hello2' * 4, wait=0.02)
self.assertEquals(dl1, dl2)
# starting a download with another fname will trigger a new download
dl3 = self.do_download('http://bar', 'foo2', 'hello you' * 4)
self.assertIsInstance(dl3, download_manager.Download)
self.assertNotEquals(dl3, dl1)
# let's wait for the downloads to finish
dl3.wait()
dl1.wait()
# now if we try to download a fname that exists, None is returned
dl4 = self.do_download('http://bar', 'foo', 'hello2' * 4, wait=0.02)
self.assertIsNone(dl4)
# downloaded files are what is expected
def content(fname):
with open(os.path.join(self.tempdir, fname)) as f:
return f.read()
self.assertEquals(content('foo'), 'hello' * 4)
self.assertEquals(content('foo2'), 'hello you' * 4)
# download instances are removed from the manager (internal test)
self.assertEquals(self.dl_manager._downloads, {})
def test_cancel(self):
dl1 = self.do_download('http://foo', 'foo', 'foo' * 50000, wait=0.02)
dl2 = self.do_download('http://foo', 'bar', 'bar' * 50000, wait=0.02)
dl3 = self.do_download('http://foo', 'foobar', 'foobar' * 4)
# let's cancel only one
def cancel_if(dl):
if os.path.basename(dl.get_dest()) == 'foo':
return True
self.dl_manager.cancel(cancel_if=cancel_if)
self.assertTrue(dl1.is_canceled())
self.assertFalse(dl2.is_canceled())
self.assertFalse(dl3.is_canceled())
# wait for dl3
dl3.wait()
# cancel everything
self.dl_manager.cancel()
self.assertTrue(dl1.is_canceled())
self.assertTrue(dl2.is_canceled())
# dl3 is not canceled since it finished before
self.assertFalse(dl3.is_canceled())
# wait for the completion of dl1 and dl2 threads
dl1.wait(raise_if_error=False)
dl2.wait(raise_if_error=False)
# at the end, only dl3 has been downloaded
self.assertEquals(os.listdir(self.tempdir), ["foobar"])
with open(os.path.join(self.tempdir, 'foobar')) as f:
self.assertEquals(f.read(), 'foobar' * 4)
# download instances are removed from the manager (internal test)
self.assertEquals(self.dl_manager._downloads, {})
|
