test path variables, comments, custom test case in own class, temp fi…

…les share remove function

datalab/datalab_session/tests/test_files/file_extended_test_case.py

@@ -0,0 +1,23 @@
+import pathlib as pl
+from hashlib import md5
+from os import path, remove, listdir
+
+from django.test import TestCase
+
+# extending the TestCase class to include a custom assertions for file operations
+class FileExtendedTestCase(TestCase):
+    def assertIsFile(self, path):
+        if not pl.Path(path).resolve().is_file():
+            raise AssertionError("File does not exist: %s" % str(path))
+
+    def assertFilesEqual(self, image_1: str, image_2: str):
+        with open(image_1, 'rb') as file_1, open(image_2, 'rb') as file_2:
+            self.assertEqual(md5(file_1.read()).hexdigest(), md5(file_2.read()).hexdigest())
+
+    def clean_test_dir(self):
+        test_files_dir = 'datalab/datalab_session/tests/test_files'
+        for file_name in listdir(test_files_dir):
+            if file_name.startswith('temp_'):
+                file_path = path.join(test_files_dir, file_name)
+                if path.isfile(file_path):
+                    remove(file_path)

datalab/datalab_session/tests/test_operations.py

@@ -1,14 +1,11 @@
 from unittest import mock
-from hashlib import md5
 import os
 
-from django.test import TestCase
-
 from datalab.datalab_session.data_operations.data_operation import BaseDataOperation
-
 from datalab.datalab_session.data_operations.rgb_stack import RGB_Stack
 from datalab.datalab_session.exceptions import ClientAlertException
 from datalab.datalab_session.data_operations.median import Median
+from datalab.datalab_session.tests.test_files.file_extended_test_case import FileExtendedTestCase
 
 wizard_description = {
             'name': 'SampleDataOperation',
@@ -38,6 +35,7 @@
                 }
             }
         }
+test_path = 'datalab/datalab_session/tests/test_files/'
 
 class SampleDataOperation(BaseDataOperation):
 
@@ -56,7 +54,7 @@ def wizard_description():
     def operate(self):
         self.set_output({'output_files': []})
 
-class TestDataOperation(TestCase):
+class TestDataOperation(FileExtendedTestCase):
     def setUp(self):
         input_data = {
             'input_files': [],
@@ -154,15 +152,14 @@ def test_set_failed(self):
         self.assertEqual(self.data_operation.get_status(), 'FAILED')
         self.assertEqual(self.data_operation.get_message(), 'Test message')
 
-class TestMedianOperation(TestCase):
-    temp_median = 'datalab/datalab_session/tests/test_files/median/temp_median.fits'
-    test_median = 'datalab/datalab_session/tests/test_files/median/median_1_2.fits'
-    test_fits_1 = 'datalab/datalab_session/tests/test_files/median/fits_1.fits.fz'
-    test_fits_2 = 'datalab/datalab_session/tests/test_files/median/fits_2.fits.fz'
+class TestMedianOperation(FileExtendedTestCase):
+    temp_median_path = f'{test_path}temp_median.fits'
+    test_median_path = f'{test_path}median/median_1_2.fits'
+    test_fits_1_path = f'{test_path}median/fits_1.fits.fz'
+    test_fits_2_path = f'{test_path}median/fits_2.fits.fz'
 
     def tearDown(self):
-        if os.path.exists(self.temp_median):
-            os.remove(self.temp_median)
+        self.clean_test_dir()
         return super().tearDown()
 
     @mock.patch('datalab.datalab_session.file_utils.tempfile.NamedTemporaryFile')
@@ -171,10 +168,14 @@ def tearDown(self):
     @mock.patch('datalab.datalab_session.data_operations.median.create_jpgs')
     def test_operate(self, mock_create_jpgs, mock_save_fits_and_thumbnails, mock_get_fits, mock_named_tempfile):
 
-        mock_get_fits.side_effect = [self.test_fits_1, self.test_fits_2]
-        mock_named_tempfile.return_value.name = self.temp_median
+        # return the test fits paths in order of the input_files instead of aws fetch
+        mock_get_fits.side_effect = [self.test_fits_1_path, self.test_fits_2_path]
+        # save temp output to a known path so we can test it
+        mock_named_tempfile.return_value.name = self.temp_median_path
+        # avoids overwriting our output
         mock_create_jpgs.return_value = ('test_path', 'test_path')
-        mock_save_fits_and_thumbnails.return_value = self.temp_median
+        # don't save to s3
+        mock_save_fits_and_thumbnails.return_value = self.temp_median_path
 
         input_data = {
             'input_files': [
@@ -189,7 +190,7 @@ def test_operate(self, mock_create_jpgs, mock_save_fits_and_thumbnails, mock_get
 
         self.assertEqual(median.get_percent_completion(), 1.0)
         self.assertTrue(os.path.exists(output[0]))
-        self.assertEqual(md5(open(self.test_median, 'rb').read()).hexdigest(), md5(open(output[0], 'rb').read()).hexdigest())
+        self.assertFilesEqual(self.test_median_path, output[0])
 
     def test_not_enough_files(self):
         input_data = {
@@ -202,16 +203,15 @@ def test_not_enough_files(self):
         with self.assertRaises(ClientAlertException):
             median.operate()
 
-class TestRGBStackOperation(TestCase):
-    temp_rgb = 'datalab/datalab_session/tests/test_files/rgb_stack/temp_rgb.fits'
-    test_rgb = 'datalab/datalab_session/tests/test_files/rgb_stack/rgb_stack.fits'
-    test_red = 'datalab/datalab_session/tests/test_files/rgb_stack/red.fits'
-    test_green = 'datalab/datalab_session/tests/test_files/rgb_stack/green.fits'
-    test_blue = 'datalab/datalab_session/tests/test_files/rgb_stack/blue.fits'
+class TestRGBStackOperation(FileExtendedTestCase):
+    temp_rgb_path = f'{test_path}temp_rgb.fits'
+    test_rgb_path = f'{test_path}rgb_stack/rgb_stack.fits'
+    test_red_path = f'{test_path}rgb_stack/red.fits'
+    test_green_path = f'{test_path}rgb_stack/green.fits'
+    test_blue_path = f'{test_path}rgb_stack/blue.fits'
 
     def tearDown(self):
-        if os.path.exists(self.temp_rgb):
-            os.remove(self.temp_rgb)
+        self.clean_test_dir()
         return super().tearDown()
 
     @mock.patch('datalab.datalab_session.data_operations.rgb_stack.save_fits_and_thumbnails')
@@ -220,10 +220,14 @@ def tearDown(self):
     @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):
 
-        mock_get_fits.side_effect = [self.test_red, self.test_green, self.test_blue]
-        mock_named_tempfile.return_value.name = self.temp_rgb
+        # return the test fits paths in order of the input_files instead of aws fetch
+        mock_get_fits.side_effect = [self.test_red_path, self.test_green_path, self.test_blue_path]
+        # save temp output to a known path so we can test
+        mock_named_tempfile.return_value.name = self.temp_rgb_path
+        # avoids overwriting our output
         mock_create_jpgs.return_value = ('test_path', 'test_path')
-        mock_save_fits_and_thumbnails.return_value = self.temp_rgb
+        # don't save to s3
+        mock_save_fits_and_thumbnails.return_value = self.temp_rgb_path
 
         input_data = {
             'red_input': [{'basename': 'red_fits', 'source': 'local'}],
@@ -237,4 +241,4 @@ def test_operate(self, mock_get_fits, mock_named_tempfile, mock_create_jpgs, moc
 
         self.assertEqual(rgb.get_percent_completion(), 1.0)
         self.assertTrue(os.path.exists(output[0]))
-        self.assertEqual(md5(open(self.test_rgb, 'rb').read()).hexdigest(), md5(open(output[0], 'rb').read()).hexdigest())
+        self.assertFilesEqual(self.test_rgb_path, output[0])

datalab/datalab_session/tests/test_utils.py

@@ -1,30 +1,18 @@
-import pathlib as pl
-from hashlib import md5
-
-from django.test import TestCase
-
 from datalab.datalab_session.file_utils import *
 from datalab.datalab_session.s3_utils import *
-
-# extending the TestCase class to include a custom assertions for file operations
-class FileExtendedTestCase(TestCase):
-    def assertIsFile(self, path):
-        if not pl.Path(path).resolve().is_file():
-            raise AssertionError("File does not exist: %s" % str(path))
-
-    def assertFilesEqual(self, image_1: str, image_2: str):
-        with open(image_1, 'rb') as file_1, open(image_2, 'rb') as file_2:
-            self.assertEqual(md5(file_1.read()).hexdigest(), md5(file_2.read()).hexdigest())
+from datalab.datalab_session.tests.test_files.file_extended_test_case import FileExtendedTestCase
 
 class FileUtilsTestClass(FileExtendedTestCase):
 
-  test_fits = 'datalab/datalab_session/tests/test_files/file_utils/fits_1.fits.fz'
-  test_tif = 'datalab/datalab_session/tests/test_files/file_utils/tif_1.tif'
-  test_small_jpg = 'datalab/datalab_session/tests/test_files/file_utils/jpg_small_1.jpg'
-  test_large_jpg = 'datalab/datalab_session/tests/test_files/file_utils/jpg_large_1.jpg'
+  util_test_path = 'datalab/datalab_session/tests/test_files/file_utils/'
+
+  test_fits_path = f'{util_test_path}fits_1.fits.fz'
+  test_tif_path = f'{util_test_path}tif_1.tif'
+  test_small_jpg_path = f'{util_test_path}jpg_small_1.jpg'
+  test_large_jpg_path = f'{util_test_path}jpg_large_1.jpg'
 
   def test_get_fits_dimensions(self):
-    fits_path = self.test_fits
+    fits_path = self.test_fits_path
     self.assertEqual(get_fits_dimensions(fits_path), (100, 100))
 
   def test_create_fits(self):
@@ -41,25 +29,25 @@ def test_create_fits(self):
     self.assertEqual(hdu[1].data.tolist(), test_2d_ndarray.tolist())
 
   def test_create_tif(self):
-    fits_path = self.test_fits
+    fits_path = self.test_fits_path
     tif_path = create_tif('create_tif_test', fits_path)
 
     # test the file was written out to a path
     self.assertIsInstance(tif_path, str)
     self.assertIsFile(tif_path)
 
-    self.assertFilesEqual(tif_path, self.test_tif)
+    self.assertFilesEqual(tif_path, self.test_tif_path)
 
   def test_create_jpgs(self):
-    fits_path = self.test_fits
+    fits_path = self.test_fits_path
     jpg_paths = create_jpgs('create_jpgs_test', fits_path)
 
     # test the files were written out to a path
     self.assertEqual(len(jpg_paths), 2)
     self.assertIsFile(jpg_paths[0])
     self.assertIsFile(jpg_paths[1])
-    self.assertFilesEqual(jpg_paths[0], self.test_large_jpg)
-    self.assertFilesEqual(jpg_paths[1], self.test_small_jpg)
+    self.assertFilesEqual(jpg_paths[0], self.test_large_jpg_path)
+    self.assertFilesEqual(jpg_paths[1], self.test_small_jpg_path)
 
   def test_stack_arrays(self):
     test_array_1 = np.zeros((10, 20))