summaryrefslogtreecommitdiffstats
path: root/third_party/python/dlmanager/tests/test_manager.py
diff options
context:
space:
mode:
Diffstat (limited to 'third_party/python/dlmanager/tests/test_manager.py')
-rw-r--r--third_party/python/dlmanager/tests/test_manager.py251
1 files changed, 251 insertions, 0 deletions
diff --git a/third_party/python/dlmanager/tests/test_manager.py b/third_party/python/dlmanager/tests/test_manager.py
new file mode 100644
index 0000000000..f0ade9021f
--- /dev/null
+++ b/third_party/python/dlmanager/tests/test_manager.py
@@ -0,0 +1,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, {})