Add more test cases for cautious-robot

src/cautiousrobot/__main__.py

@@ -145,10 +145,17 @@ def download_images(data, img_dir, log_filepath, error_log_filepath, filename =
                         if os.path.exists(downsample_dir_path) != True:
                             os.makedirs(downsample_dir_path, exist_ok=False)
                         # Downsample & save image
-                        byte_data = io.BytesIO(response.content)
-                        img = Image.open(byte_data)
-                        #img.save(dest_path)
-                        img.resize((downsample, downsample)).save(downsample_dir_path + "/" + image_name)
+                        try:
+                            img = Image.open(f"{image_dir_path}/{image_name}")
+                            img.resize((downsample, downsample)).save(downsample_dir_path + "/" + image_name)
+                        except Exception as e:
+                            print(e)
+                            log_errors = log_response(log_errors,
+                                        index = i,
+                                        image = "downsized_" + image_name,
+                                        url = url,
+                                        response_code = str(e))
+                            update_log(log = log_errors, index = i, filepath = error_log_filepath)
 
                 # check for too many requests
                 elif response.status_code in REDO_CODE_LIST:

tests/test_buddycheck.py

@@ -46,8 +46,9 @@ def test_merge_on_checksum(self):
 
         merged_df = self.buddy_check_filename.merge_on_checksum(source_df, checksum_df, 'checksum')
         expected_df = pd.DataFrame({
-            'filename': ['image1.jpg', 'image2.jpg', 'image3.jpg'],
+            'filename_x': ['image1.jpg', 'image2.jpg', 'image3.jpg'],
             'checksum': ['abc123', 'def456', 'ghi789'],
+            'filename_y': ['image1.jpg', 'image2.jpg', 'image3.jpg'],
             'md5': ['abc123', 'def456', 'ghi789']
         })
         pd.testing.assert_frame_equal(merged_df, expected_df)

tests/test_download_images.py

@@ -9,26 +9,50 @@
 import requests
 import base64
 from cautiousrobot.__main__ import download_images, main
+from http.server import HTTPServer, SimpleHTTPRequestHandler
+import threading
+
+TESTDATA_DIR = os.path.join(os.path.dirname(__file__), 'testdata')
+
+class CustomHTTPRequestHandler(SimpleHTTPRequestHandler):
+    def translate_path(self, path):
+        return os.path.join(TESTDATA_DIR, os.path.relpath(path, '/'))
 
 class TestDownload(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.httpd = HTTPServer(('localhost', 9201), CustomHTTPRequestHandler)
+        cls.server_thread = threading.Thread(target=cls.httpd.serve_forever)
+        cls.server_thread.start()
+        print(f"Serving {TESTDATA_DIR} on http://localhost:9201")
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.httpd.shutdown()
+        cls.server_thread.join()
+
+
     def setUp(self):
         self.DUMMY_DATA = pd.DataFrame(data={
-            "filename": ["test_file1", "test_file2"],
-            "file_url": ["http://test_url1.com/image.jpg", "http://test_url2.com/image.jpg"],
+            "filename": ["test_file1.jpg", "test_file2.jpg"],
+            "file_url": ["http://localhost:9201/images/image1.jpg", "http://localhost:9201/images/image2.png"],
             "subfolder": ["test_subfolder1", "test_subfolder2"]
         })
         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):
@@ -54,7 +78,22 @@ def test_successful_download(self, get_mock):
             self.assertTrue(os.path.isfile(f"{self.IMG_DIR}/{filename}"))
 
     @patch('requests.get')
-    def test_success_after_retries(self, get_mock):
+    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')
+    @patch('time.sleep', return_value=None)
+    def test_success_after_retries(self,sleep_mock, get_mock):
         retry_status_codes = [429, 500, 502, 503, 504]
         for status_code in retry_status_codes:
             with self.subTest(status_code=status_code):
@@ -63,39 +102,51 @@ 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):
+    @patch('time.sleep', return_value=None)
+    def test_failure_after_retries(self, sleep_mock,get_mock):
         retry_status_codes = [429, 500, 502, 503, 504]
         for status_code in retry_status_codes:
             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}"))
-
-    @patch('cautiousrobot.__main__.requests.get')
-    def test_downsampled_image_creation(self,get_mock):
-        mock_response = MagicMock()
-        mock_response.content = b'dummy_image_data'
-        get_mock.return_value = mock_response
+            for filename in self.DUMMY_DATA['filename']:
+                self.assertFalse(os.path.isfile(f"{self.IMG_DIR}/{filename}"))
 
+    def test_downsampled_image_creation(self):
         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}"))    
+            downsampled_path = os.path.join(self.DOWNSAMPLE_DIR, filename)
+            print(f"Checking existence of downsampled image: {downsampled_path}")
+            self.assertTrue(os.path.isfile(f"{self.DOWNSAMPLE_DIR}/{filename}") or
+                            os.path.isfile(f"{self.DOWNSAMPLE_DIR}/{filename}"))   
+
+    def test_downsampled_image_creation_with_subfolder(self):
+        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):