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