Source code
Revision control
Copy as Markdown
Other Tools
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')
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_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):
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.
self.assertEquals(dl1, dl2)
# starting a download with another fname will trigger a new download
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
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):
# 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, {})