bgneal@1225: import unittest
bgneal@1225: from urlparse import urlparse, ParseResult
bgneal@1225: 
bgneal@1225: from mock import call, patch, mock_open, Mock
bgneal@1225: 
bgneal@1225: from core.download import download_file
bgneal@1225: 
bgneal@1225: 
bgneal@1225: class DownloadFileTestCase(unittest.TestCase):
bgneal@1225:     """Unit tests for the download_file function."""
bgneal@1225: 
bgneal@1225:     @patch('core.download.requests')
bgneal@1225:     def test_get_throws_exception(self, requests_mock):
bgneal@1225:         requests_mock.get.side_effect = RuntimeError('ope')
bgneal@1225:         r = None
bgneal@1225:         try:
bgneal@1225:             r = download_file('url', path='path', timeout=5.0)
bgneal@1225:         except RuntimeError:
bgneal@1225:             pass
bgneal@1225:         except Exception as e:
bgneal@1225:             self.fail('Unexpected exception {}'.format(e))
bgneal@1225: 
bgneal@1225:         self.assertIsNone(r)
bgneal@1225:         self.assertEqual(requests_mock.get.mock_calls, [
bgneal@1225:             call('url', stream=True, timeout=5.0),
bgneal@1225:         ])
bgneal@1225: 
bgneal@1225:     @patch('core.download.requests')
bgneal@1225:     def test_get_returns_non_success(self, requests_mock):
bgneal@1225:         response_mock = Mock()
bgneal@1225:         response_mock.status_code = 404
bgneal@1225:         requests_mock.get.return_value = response_mock
bgneal@1225:         r = download_file('url', path='path', timeout=5.0)
bgneal@1225: 
bgneal@1225:         self.assertIsNone(r)
bgneal@1225:         self.assertEqual(requests_mock.get.mock_calls, [
bgneal@1225:             call('url', stream=True, timeout=5.0),
bgneal@1225:         ])
bgneal@1225: 
bgneal@1225:     @patch('core.download.os.stat')
bgneal@1225:     @patch('core.download.shutil.copyfileobj')
bgneal@1225:     @patch('core.download.requests')
bgneal@1225:     def test_happy_path_with_path_argument(self, requests_mock, copy_mock,
bgneal@1225:                                            stat_mock):
bgneal@1225:         response_mock = Mock()
bgneal@1225:         response_mock.status_code = 200
bgneal@1225:         requests_mock.get.return_value = response_mock
bgneal@1225: 
bgneal@1225:         stat_mock.st_size = 512
bgneal@1225: 
bgneal@1225:         open_mock = mock_open()
bgneal@1225:         with patch('__builtin__.open', open_mock):
bgneal@1225:             r = download_file('url', path='path', timeout=5.0)
bgneal@1225: 
bgneal@1225:         self.assertEqual(r, 'path')
bgneal@1225:         self.assertEqual(requests_mock.get.mock_calls, [
bgneal@1225:             call('url', stream=True, timeout=5.0),
bgneal@1225:         ])
bgneal@1225:         self.assertTrue(response_mock.raw.decode_content)
bgneal@1225:         self.assertEqual(copy_mock.mock_calls, [
bgneal@1225:             call(response_mock.raw, open_mock.return_value),
bgneal@1225:         ])
bgneal@1225: 
bgneal@1225:     @patch('core.download.os.remove')
bgneal@1225:     @patch('core.download.shutil.copyfileobj')
bgneal@1225:     @patch('core.download.requests')
bgneal@1225:     def test_copyfileobj_raises(self, requests_mock, copy_mock, remove_mock):
bgneal@1225:         response_mock = Mock()
bgneal@1225:         response_mock.status_code = 200
bgneal@1225:         requests_mock.get.return_value = response_mock
bgneal@1225: 
bgneal@1225:         copy_mock.side_effect = IOError
bgneal@1225: 
bgneal@1225:         open_mock = mock_open()
bgneal@1225:         with patch('__builtin__.open', open_mock):
bgneal@1225:             try:
bgneal@1225:                 download_file('url', path='path', timeout=5.0)
bgneal@1225:             except IOError:
bgneal@1225:                 pass
bgneal@1225:             else:
bgneal@1225:                 self.fail('Should have thrown')
bgneal@1225: 
bgneal@1225:         self.assertEqual(requests_mock.get.mock_calls, [
bgneal@1225:             call('url', stream=True, timeout=5.0),
bgneal@1225:         ])
bgneal@1225:         self.assertTrue(response_mock.raw.decode_content)
bgneal@1225:         self.assertEqual(copy_mock.mock_calls, [
bgneal@1225:             call(response_mock.raw, open_mock.return_value),
bgneal@1225:         ])
bgneal@1225:         self.assertEqual(remove_mock.mock_calls, [
bgneal@1225:             call('path'),
bgneal@1225:         ])
bgneal@1225: 
bgneal@1225: 
bgneal@1225:     @patch('core.download.os')
bgneal@1225:     @patch('core.download.tempfile')
bgneal@1225:     @patch('core.download.mimetypes')
bgneal@1225:     @patch('core.download.os.remove')
bgneal@1225:     @patch('core.download.shutil.copyfileobj')
bgneal@1225:     @patch('core.download.requests')
bgneal@1225:     def test_happy_path_with_suffix(self, requests_mock, copy_mock, remove_mock,
bgneal@1225:                                     mime_mock, tempfile_mock, os_mock):
bgneal@1225:         response_mock = Mock()
bgneal@1225:         response_mock.status_code = 200
bgneal@1225:         requests_mock.get.return_value = response_mock
bgneal@1225: 
bgneal@1225:         response_mock.headers.get.return_value = 'image/jpeg'
bgneal@1225:         mime_mock.guess_extension.return_value = '.jpe'
bgneal@1225:         fd_mock = Mock()
bgneal@1225:         tempfile_mock.mkstemp.return_value = (fd_mock, 'temp-path')
bgneal@1225: 
bgneal@1225:         open_mock = mock_open()
bgneal@1225:         with patch('__builtin__.open', open_mock):
bgneal@1225:             download_file('url', timeout=5.0)
bgneal@1225: 
bgneal@1225:         self.assertEqual(requests_mock.get.mock_calls, [
bgneal@1225:             call('url', stream=True, timeout=5.0),
bgneal@1225:             call().headers.get('content-type'),
bgneal@1225:         ])
bgneal@1225:         self.assertEqual(response_mock.headers.get.mock_calls, [
bgneal@1225:             call('content-type'),
bgneal@1225:         ])
bgneal@1225:         self.assertEqual(mime_mock.guess_extension.mock_calls, [
bgneal@1225:             call('image/jpeg'),
bgneal@1225:         ])
bgneal@1225:         self.assertEqual(tempfile_mock.mkstemp.mock_calls, [
bgneal@1225:             call(suffix='.jpg'),
bgneal@1225:         ])
bgneal@1225:         self.assertEqual(os_mock.close.mock_calls, [
bgneal@1225:             call(fd_mock),
bgneal@1225:         ])
bgneal@1225:         self.assertTrue(response_mock.raw.decode_content)
bgneal@1225:         self.assertEqual(copy_mock.mock_calls, [
bgneal@1225:             call(response_mock.raw, open_mock.return_value),
bgneal@1225:         ])
bgneal@1225: 
bgneal@1225: 
bgneal@1225:     @patch('core.download.urlparse')
bgneal@1225:     @patch('core.download.os')
bgneal@1225:     @patch('core.download.tempfile')
bgneal@1225:     @patch('core.download.mimetypes')
bgneal@1225:     @patch('core.download.os.remove')
bgneal@1225:     @patch('core.download.shutil.copyfileobj')
bgneal@1225:     @patch('core.download.requests')
bgneal@1225:     def test_happy_path_with_no_suffix(self, requests_mock, copy_mock,
bgneal@1225:                                        remove_mock, mime_mock, tempfile_mock,
bgneal@1225:                                        os_mock, urlparse_mock):
bgneal@1225:         response_mock = Mock()
bgneal@1225:         response_mock.status_code = 200
bgneal@1225:         requests_mock.get.return_value = response_mock
bgneal@1225: 
bgneal@1225:         response_mock.headers.get.return_value = None
bgneal@1225:         urlparse_mock.return_value = ParseResult(
bgneal@1225:                 scheme='https', netloc='www.example.com', path='/something.txt',
bgneal@1225:                 params='', query='', fragment='')
bgneal@1225:         os_mock.path.splitext.return_value = ('/something', '.txt')
bgneal@1225:         fd_mock = Mock()
bgneal@1225:         tempfile_mock.mkstemp.return_value = (fd_mock, 'temp-path')
bgneal@1225: 
bgneal@1225:         open_mock = mock_open()
bgneal@1225:         with patch('__builtin__.open', open_mock):
bgneal@1225:             download_file('url', timeout=5.0)
bgneal@1225: 
bgneal@1225:         self.assertEqual(requests_mock.get.mock_calls, [
bgneal@1225:             call('url', stream=True, timeout=5.0),
bgneal@1225:             call().headers.get('content-type'),
bgneal@1225:         ])
bgneal@1225:         self.assertEqual(response_mock.headers.get.mock_calls, [
bgneal@1225:             call('content-type'),
bgneal@1225:         ])
bgneal@1225:         self.assertEqual(urlparse_mock.mock_calls, [
bgneal@1225:             call('url'),
bgneal@1225:         ])
bgneal@1225:         self.assertEqual(mime_mock.guess_extension.mock_calls, [])
bgneal@1225:         self.assertEqual(os_mock.path.splitext.mock_calls, [
bgneal@1225:             call('/something.txt'),
bgneal@1225:         ])
bgneal@1225:         self.assertEqual(tempfile_mock.mkstemp.mock_calls, [
bgneal@1225:             call(suffix='.txt'),
bgneal@1225:         ])
bgneal@1225:         self.assertEqual(os_mock.close.mock_calls, [
bgneal@1225:             call(fd_mock),
bgneal@1225:         ])
bgneal@1225:         self.assertTrue(response_mock.raw.decode_content)
bgneal@1225:         self.assertEqual(copy_mock.mock_calls, [
bgneal@1225:             call(response_mock.raw, open_mock.return_value),
bgneal@1225:         ])