From e92f93a0563a492dbf870ee0b5114e09e2e5056f Mon Sep 17 00:00:00 2001 From: zoeduan Date: Mon, 29 Jul 2024 13:41:21 -0400 Subject: [PATCH] 1. Update test_success_after_retries and test_failure_after_retries 2. Add testcases test_successful_download_with_subfolder and test_downsampled_image_creation_with_subfolder --- tests/test_download_images.py | 59 ++++++++++++++++++++++++++++------- 1 file changed, 47 insertions(+), 12 deletions(-) diff --git a/tests/test_download_images.py b/tests/test_download_images.py index 45406c5..d3f9472 100644 --- a/tests/test_download_images.py +++ b/tests/test_download_images.py @@ -20,15 +20,17 @@ def setUp(self): self.IMG_DIR = "test_dir" self.LOG_FILEPATH = "test_log_path.jsonl" self.ERROR_LOG_FILEPATH = "test_error_log_path.jsonl" - self.DOWNSAMPLE_PATH = "test_downsample_dir" + self.DOWNSAMPLE_DIR = self.IMG_DIR + "_downsized" self.DOWNSAMPLE_SIZE = 100 os.makedirs(self.IMG_DIR, exist_ok=True) - os.makedirs(self.DOWNSAMPLE_PATH, exist_ok=True) + os.makedirs(self.DOWNSAMPLE_DIR, exist_ok=True) + for subfolder in self.DUMMY_DATA["subfolder"]: + os.makedirs(os.path.join(self.DOWNSAMPLE_DIR, subfolder), exist_ok=True) def tearDown(self): shutil.rmtree(self.IMG_DIR, ignore_errors=True) - shutil.rmtree(self.DOWNSAMPLE_PATH, ignore_errors=True) + shutil.rmtree(self.DOWNSAMPLE_DIR, ignore_errors=True) if os.path.exists(self.LOG_FILEPATH): os.remove(self.LOG_FILEPATH) if os.path.exists(self.ERROR_LOG_FILEPATH): @@ -53,6 +55,20 @@ def test_successful_download(self, get_mock): for filename in self.DUMMY_DATA['filename']: self.assertTrue(os.path.isfile(f"{self.IMG_DIR}/{filename}")) + @patch('requests.get') + def test_successful_download_with_subfolder(self, get_mock): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raw = BytesIO(b"fake_image_data") + get_mock.return_value = mock_response + + download_images(self.DUMMY_DATA, self.IMG_DIR, self.LOG_FILEPATH, self.ERROR_LOG_FILEPATH, subfolders="subfolder") + + for i, filename in enumerate(self.DUMMY_DATA['filename']): + subfolder = self.DUMMY_DATA['subfolder'][i] + self.assertTrue(os.path.isfile(f"{self.IMG_DIR}/{subfolder}/{filename}")) + + @patch('requests.get') def test_success_after_retries(self, get_mock): retry_status_codes = [429, 500, 502, 503, 504] @@ -63,12 +79,15 @@ def test_success_after_retries(self, get_mock): mock_response_success = MagicMock() mock_response_success.status_code = 200 mock_response_success.raw = BytesIO(b"fake_image_data") - get_mock.side_effect = [mock_response_retry] * 2 + [mock_response_success] + get_mock.side_effect = [ + mock_response_retry, mock_response_retry, mock_response_success, # For test_file1.jpg + mock_response_retry, mock_response_retry, mock_response_success # For test_file2.jpg + ] download_images(self.DUMMY_DATA, self.IMG_DIR, self.LOG_FILEPATH, self.ERROR_LOG_FILEPATH, retry=3) - for filename in self.DUMMY_DATA['filename']: - self.assertTrue(os.path.isfile(f"{self.IMG_DIR}/{filename}")) + for filename in self.DUMMY_DATA['filename']: + self.assertTrue(os.path.isfile(f"{self.IMG_DIR}/{filename}")) @patch('requests.get') def test_failure_after_retries(self, get_mock): @@ -77,12 +96,15 @@ def test_failure_after_retries(self, get_mock): with self.subTest(status_code=status_code): mock_response_retry = MagicMock() mock_response_retry.status_code = status_code - get_mock.side_effect = [mock_response_retry] * 5 + get_mock.side_effect = [ + mock_response_retry, mock_response_retry, mock_response_retry, mock_response_retry, mock_response_retry, # For test_file1.jpg + mock_response_retry, mock_response_retry, mock_response_retry, mock_response_retry, mock_response_retry # For test_file2.jpg + ] download_images(self.DUMMY_DATA, self.IMG_DIR, self.LOG_FILEPATH, self.ERROR_LOG_FILEPATH, retry=5) - for filename in self.DUMMY_DATA['filename']: - self.assertFalse(os.path.isfile(f"{self.IMG_DIR}/{filename}")) + for filename in self.DUMMY_DATA['filename']: + self.assertFalse(os.path.isfile(f"{self.IMG_DIR}/{filename}")) @patch('cautiousrobot.__main__.requests.get') def test_downsampled_image_creation(self,get_mock): @@ -91,11 +113,24 @@ def test_downsampled_image_creation(self,get_mock): get_mock.return_value = mock_response download_images(self.DUMMY_DATA, self.IMG_DIR, self.LOG_FILEPATH, self.ERROR_LOG_FILEPATH, - downsample_path=self.DOWNSAMPLE_PATH, downsample=self.DOWNSAMPLE_SIZE) + downsample_path=self.DOWNSAMPLE_DIR, downsample=self.DOWNSAMPLE_SIZE) for filename in self.DUMMY_DATA['filename']: - self.assertTrue(os.path.isfile(f"{self.DOWNSAMPLE_PATH}/test_subfolder1/{filename}") or - os.path.isfile(f"{self.DOWNSAMPLE_PATH}/test_subfolder2/{filename}")) + self.assertTrue(os.path.isfile(f"{self.DOWNSAMPLE_DIR}/{filename}") or + os.path.isfile(f"{self.DOWNSAMPLE_DIR}/{filename}")) + + @patch('cautiousrobot.__main__.requests.get') + def test_downsampled_image_creation_with_subfolder(self, get_mock): + mock_response = MagicMock() + mock_response.content = b'dummy_image_data' + get_mock.return_value = mock_response + + download_images(self.DUMMY_DATA, self.IMG_DIR, self.LOG_FILEPATH, self.ERROR_LOG_FILEPATH, + downsample_path=self.DOWNSAMPLE_DIR, downsample=self.DOWNSAMPLE_SIZE, subfolders="subfolder") + + for i, filename in enumerate(self.DUMMY_DATA['filename']): + subfolder = self.DUMMY_DATA['subfolder'][i] + self.assertTrue(os.path.isfile(f"{self.DOWNSAMPLE_DIR}/{subfolder}/{filename}")) @patch('requests.get') def test_logging(self, get_mock):