fix tests for create_output change, overwrite temp files

LCOGT · Oct 2, 2024 · e0411d1 · e0411d1
1 parent 26ed7f5
commit e0411d1
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 12 deletions.
diff --git a/datalab/datalab_session/analysis/get_tif.py b/datalab/datalab_session/analysis/get_tif.py
@@ -1,5 +1,5 @@
-from datalab.datalab_session.file_utils import create_tif, get_fits
-from datalab.datalab_session.s3_utils import key_exists, add_file_to_bucket, get_s3_url
+from datalab.datalab_session.file_utils import create_tif
+from datalab.datalab_session.s3_utils import key_exists, add_file_to_bucket, get_s3_url, get_fits
 
 def get_tif(input: dict):
   """

diff --git a/datalab/datalab_session/data_operations/rgb_stack.py b/datalab/datalab_session/data_operations/rgb_stack.py
@@ -5,7 +5,8 @@
 
 from datalab.datalab_session.data_operations.data_operation import BaseDataOperation
 from datalab.datalab_session.exceptions import ClientAlertException
-from datalab.datalab_session.file_utils import create_output, get_fits, crop_arrays, create_jpgs
+from datalab.datalab_session.file_utils import create_output, crop_arrays, create_jpgs
+from datalab.datalab_session.s3_utils import get_fits
 
 log = logging.getLogger()
 log.setLevel(logging.INFO)

diff --git a/datalab/datalab_session/file_utils.py b/datalab/datalab_session/file_utils.py
@@ -39,7 +39,7 @@ def create_fits(key: str, image_arr: np.ndarray) -> str:
 
   hdu_list = fits.HDUList([primary_hdu, image_hdu])
   fits_path = tempfile.NamedTemporaryFile(suffix=f'{key}.fits').name
-  hdu_list.writeto(fits_path)
+  hdu_list.writeto(fits_path, overwrite=True)
 
   return fits_path
 

diff --git a/datalab/datalab_session/tests/test_operations.py b/datalab/datalab_session/tests/test_operations.py
@@ -56,7 +56,7 @@ def wizard_description():
         return wizard_description
 
     def operate(self):
-        self.set_output({'output_files': []})
+        self.set_output([])
 
 
 class TestDataOperation(FileExtendedTestCase):
@@ -147,7 +147,7 @@ def test_generate_cache_key(self):
         self.assertEqual(self.data_operation.generate_cache_key(), pregenerated_cache_key)
 
     def test_set_get_output(self):
-        self.data_operation.set_output({'output_files': []})
+        self.data_operation.set_output([])
         self.assertEqual(self.data_operation.get_operation_progress(), 1.0)
         self.assertEqual(self.data_operation.get_status(), 'COMPLETED')
         self.assertEqual(self.data_operation.get_output(), {'output_files': []})
@@ -170,8 +170,8 @@ def tearDown(self):
 
     @mock.patch('datalab.datalab_session.file_utils.tempfile.NamedTemporaryFile')
     @mock.patch('datalab.datalab_session.data_operations.data_operation.get_fits')
-    @mock.patch('datalab.datalab_session.data_operations.median.save_fits_and_thumbnails')
-    @mock.patch('datalab.datalab_session.data_operations.median.create_jpgs')
+    @mock.patch('datalab.datalab_session.file_utils.save_fits_and_thumbnails')
+    @mock.patch('datalab.datalab_session.file_utils.create_jpgs')
     def test_operate(self, mock_create_jpgs, mock_save_fits_and_thumbnails, mock_get_fits, mock_named_tempfile):
 
         # return the test fits paths in order of the input_files instead of aws fetch
@@ -221,8 +221,8 @@ def tearDown(self):
         self.clean_test_dir()
         return super().tearDown()
 
-    @mock.patch('datalab.datalab_session.data_operations.rgb_stack.save_fits_and_thumbnails')
-    @mock.patch('datalab.datalab_session.data_operations.rgb_stack.create_jpgs')
+    @mock.patch('datalab.datalab_session.file_utils.save_fits_and_thumbnails')
+    @mock.patch('datalab.datalab_session.file_utils.create_jpgs')
     @mock.patch('datalab.datalab_session.file_utils.tempfile.NamedTemporaryFile')
     @mock.patch('datalab.datalab_session.data_operations.rgb_stack.get_fits')
     def test_operate(self, mock_get_fits, mock_named_tempfile, mock_create_jpgs, mock_save_fits_and_thumbnails):
@@ -266,8 +266,8 @@ def tearDown(self):
 
     @mock.patch('datalab.datalab_session.file_utils.tempfile.NamedTemporaryFile')
     @mock.patch('datalab.datalab_session.data_operations.data_operation.get_fits')
-    @mock.patch('datalab.datalab_session.data_operations.stacking.save_fits_and_thumbnails')
-    @mock.patch('datalab.datalab_session.data_operations.stacking.create_jpgs')
+    @mock.patch('datalab.datalab_session.file_utils.save_fits_and_thumbnails')
+    @mock.patch('datalab.datalab_session.file_utils.create_jpgs')
     def test_operate(self, mock_create_jpgs, mock_save_fits_and_thumbnails, mock_get_fits, mock_named_tempfile):
 
         # Create a negative images using numpy