better error handling in line profile, analysis integration tests, an…

…alysis testing data
LCOGT · Aug 28, 2024 · 656f93e · 656f93e
1 parent feee94c
commit 656f93e
Show file tree

Hide file tree

Showing 5 changed files with 122 additions and 28 deletions.
diff --git a/datalab/datalab_session/analysis/line_profile.py b/datalab/datalab_session/analysis/line_profile.py
@@ -30,20 +30,26 @@ def line_profile(input: dict):
   try: 
     wcs = WCS(sci_hdu.header)
 
+    if(wcs.get_axis_types()[0].get('coordinate_type') == None):
+      raise WcsError("No valid WCS solution")
+
     start_sky_coord = wcs.pixel_to_world(x_points[0], y_points[0])
     end_sky_coord = wcs.pixel_to_world(x_points[1], y_points[1])
 
     arcsec_angle = start_sky_coord.separation(end_sky_coord).arcsecond
 
     start_coords = [start_sky_coord.ra.deg, start_sky_coord.dec.deg]
     end_coords = [end_sky_coord.ra.deg, end_sky_coord.dec.deg]
-
-  except WcsError as e:
-    error = f'{input["basename"]} does not have a valid WCS header'
-    print(error, e)
-    # if theres no valid WCS solution then we default to using pixscale to calculate the angle, and no coordinates
+  except WcsError:
+    # no valid WCS solution
     start_coords = None
     end_coords = None
-    arcsec_angle = len(line_profile) * sci_hdu.header["PIXSCALE"]
+
+    try:
+      # attempt using pixscale to calculate the angle
+      arcsec_angle = len(line_profile) * sci_hdu.header["PIXSCALE"]
+    except KeyError as e:
+      # no valid WCS solution, and no pixscale
+      arcsec_angle = None
 
   return {"line_profile": line_profile, "arcsec": arcsec_angle, "start_coords": start_coords, "end_coords": end_coords}
diff --git a/datalab/datalab_session/tests/test_analysis.py b/datalab/datalab_session/tests/test_analysis.py
@@ -1,28 +1,48 @@
 from unittest import mock
+import json
 
 from django.test import TestCase
+from numpy.testing import assert_almost_equal
 
-class TestAnalysis(TestCase):
-    def setUp(self):
-        pass
-
-    def test_get_tif(self):
-        # TODO use a test fits file to create a tif file
-
-        # TODO mock the get_s3_url call
-        # TODO mock the get_fits call
-        # TODO mock the add_file_to_bucket call
+from datalab.datalab_session.analysis import line_profile, source_catalog
 
-        # TODO assert tif file exists
-        # TODO assert tif file exists pregenerated tif file
-        pass
-
-    def test_line_profile(self):
-        # TODO use a test fits file that doesn't have a WCS header
-
-        # TODO assert that we throw a WCS error if the header is invalid
-        pass
+class TestAnalysis(TestCase):
+    analysis_test_path = 'datalab/datalab_session/tests/test_files/analysis/'
+    analysis_fits_1_path = f'{analysis_test_path}fits_1.fits.fz'
 
-    def test_source_catalog(self):
-        # TODO use a test fits file 
-        pass
+    def setUp(self):
+        with open(f'{self.analysis_test_path}test_line_profile.json') as f:
+            self.test_line_profile_data = json.load(f)['test_line_profile']
+
+        with open(f'{self.analysis_test_path}test_source_catalog.json') as f:
+            self.test_source_catalog_data = json.load(f)['test_source_catalog']
+
+    @mock.patch('datalab.datalab_session.file_utils.get_fits')
+    def test_line_profile(self, mock_get_fits):
+
+        mock_get_fits.return_value = self.analysis_fits_1_path
+
+        output = line_profile.line_profile({
+            'basename': 'fits_1',
+            'height': 100,
+            'width': 100,
+            'x1': 25,
+            'y1': 25,
+            'x2': 75,
+            'y2': 75
+        })
+
+        assert_almost_equal(output.get('line_profile').tolist(), self.test_line_profile_data, decimal=3)
+
+    @mock.patch('datalab.datalab_session.file_utils.get_fits')
+    def test_source_catalog(self, mock_get_fits):
+
+        mock_get_fits.return_value = self.analysis_fits_1_path
+
+        output = source_catalog.source_catalog({
+            'basename': 'fits_1',
+            'height': 100,
+            'width': 100,
+            })
+
+        self.assertEqual(output, self.test_source_catalog_data)
diff --git a/datalab/datalab_session/tests/test_files/analysis/fits_1.fits.fz b/datalab/datalab_session/tests/test_files/analysis/fits_1.fits.fz
diff --git a/datalab/datalab_session/tests/test_files/analysis/test_line_profile.json b/datalab/datalab_session/tests/test_files/analysis/test_line_profile.json
@@ -0,0 +1,14 @@
+{
+  "test_line_profile": [135.66614, 187.97342, 186.27275, 152.81543, 200.08119, 190.92934,
+       190.35522, 200.84515, 185.5646 , 182.38306, 182.62105, 195.83347,
+       181.14099, 170.32465, 193.30614, 189.06567, 172.27455, 194.88974,
+       184.23103, 183.86188, 177.11359, 184.32747, 185.9386 , 177.15134,
+       171.93048, 183.24039, 188.6636 , 171.44247, 175.03983, 169.98991,
+       201.46959, 173.0296 , 172.73636, 169.93402, 170.76671, 177.4004 ,
+       182.55852, 190.3085 , 181.39314, 181.369  , 171.61351, 158.37466,
+       178.958  , 168.84396, 170.62454, 186.83702, 173.42389, 165.36671,
+       186.99101, 181.13121, 153.41772, 185.963  , 165.45833, 175.77798,
+       180.02434, 177.85484, 182.0603 , 178.34627, 178.07858, 180.45233,
+       170.2989 , 176.22253, 169.75783, 171.57878, 180.5748 , 180.7788 ,
+       185.37752, 180.1012 , 157.70186, 181.05095, 168.38719, 159.55258]
+}
diff --git a/datalab/datalab_session/tests/test_files/analysis/test_source_catalog.json b/datalab/datalab_session/tests/test_files/analysis/test_source_catalog.json
@@ -0,0 +1,54 @@
+{
+  "test_source_catalog": [
+    {"x": 792, "y": 623, "flux": 3164455, "ra": "260.9333", "dec": "-33.0225"},
+    {"x": 919, "y": 120, "flux": 807813, "ra": "260.9667", "dec": "-32.9137"}, 
+    {"x": 996, "y": 627, "flux": 383099, "ra": "260.9861", "dec": "-33.0235"}, 
+    {"x": 997, "y": 641, "flux": 335359, "ra": "260.9862", "dec": "-33.0265"}, 
+    {"x": 587, "y": 803, "flux": 304833, "ra": "260.8802", "dec": "-33.0612"}, 
+    {"x": 1015, "y": 554, "flux": 291847, "ra": "260.9910", "dec": "-33.0076"}, 
+    {"x": 490, "y": 568, "flux": 226494, "ra": "260.8555", "dec": "-33.0101"}, 
+    {"x": 526, "y": 147, "flux": 209307, "ra": "260.8654", "dec": "-32.9191"}, 
+    {"x": 740, "y": 301, "flux": 177543, "ra": "260.9203", "dec": "-32.9526"}, 
+    {"x": 334, "y": 231, "flux": 168666, "ra": "260.8157", "dec": "-32.9370"}, 
+    {"x": 843, "y": 277, "flux": 132858, "ra": "260.9468", "dec": "-32.9476"}, 
+    {"x": 776, "y": 805, "flux": 115180, "ra": "260.9290", "dec": "-33.0618"}, 
+    {"x": 13, "y": 341, "flux": 101317, "ra": "260.7327", "dec": "-32.9604"}, 
+    {"x": 384, "y": 332, "flux": 93054, "ra": "260.8284", "dec": "-32.9590"}, 
+    {"x": 451, "y": 307, "flux": 82848, "ra": "260.8457", "dec": "-32.9536"}, 
+    {"x": 830, "y": 815, "flux": 76200, "ra": "260.9427", "dec": "-33.0641"}, 
+    {"x": 846, "y": 138, "flux": 70788, "ra": "260.9478", "dec": "-32.9176"}, 
+    {"x": 569, "y": 179, "flux": 69761, "ra": "260.8763", "dec": "-32.9261"}, 
+    {"x": 818, "y": 146, "flux": 69329, "ra": "260.9406", "dec": "-32.9191"}, 
+    {"x": 843, "y": 912, "flux": 63748, "ra": "260.9462", "dec": "-33.0851"}, 
+    {"x": 395, "y": 151, "flux": 60681, "ra": "260.8315", "dec": "-32.9198"}, 
+    {"x": 542, "y": 122, "flux": 60653, "ra": "260.8696", "dec": "-32.9136"}, 
+    {"x": 214, "y": 390, "flux": 60057, "ra": "260.7845", "dec": "-32.9712"}, 
+    {"x": 188, "y": 691, "flux": 59835, "ra": "260.7774", "dec": "-33.0363"}, 
+    {"x": 173, "y": 438, "flux": 59768, "ra": "260.7738", "dec": "-32.9816"}, 
+    {"x": 1022, "y": 188, "flux": 56811, "ra": "260.9931", "dec": "-32.9285"}, 
+    {"x": 431, "y": 936, "flux": 56575, "ra": "260.8396", "dec": "-33.0898"}, 
+    {"x": 333, "y": 112, "flux": 51136, "ra": "260.8157", "dec": "-32.9112"}, 
+    {"x": 211, "y": 958, "flux": 50722, "ra": "260.7828", "dec": "-33.0941"}, 
+    {"x": 664, "y": 535, "flux": 45352, "ra": "260.9004", "dec": "-33.0033"}, 
+    {"x": 989, "y": 975, "flux": 44414, "ra": "260.9838", "dec": "-33.0988"}, 
+    {"x": 1003, "y": 463, "flux": 38238, "ra": "260.9878", "dec": "-32.9879"}, 
+    {"x": 336, "y": 242, "flux": 37970, "ra": "260.8161", "dec": "-32.9393"}, 
+    {"x": 294, "y": 565, "flux": 37707, "ra": "260.8050", "dec": "-33.0093"}, 
+    {"x": 535, "y": 916, "flux": 37575, "ra": "260.8666", "dec": "-33.0854"}, 
+    {"x": 66, "y": 415, "flux": 37365, "ra": "260.7463", "dec": "-32.9765"}, 
+    {"x": 533, "y": 91, "flux": 36539, "ra": "260.8673", "dec": "-32.9069"}, 
+    {"x": 879, "y": 505, "flux": 35138, "ra": "260.9559", "dec": "-32.9970"}, 
+    {"x": 293, "y": 638, "flux": 34703, "ra": "260.8045", "dec": "-33.0250"}, 
+    {"x": 444, "y": 635, "flux": 33406, "ra": "260.8434", "dec": "-33.0245"}, 
+    {"x": 400, "y": 82, "flux": 33366, "ra": "260.8330", "dec": "-32.9048"}, 
+    {"x": 465, "y": 59, "flux": 32920, "ra": "260.8497", "dec": "-32.8998"}, 
+    {"x": 874, "y": 50, "flux": 32280, "ra": "260.9553", "dec": "-32.8984"}, 
+    {"x": 1007, "y": 416, "flux": 29422, "ra": "260.9892", "dec": "-32.9779"}, 
+    {"x": 281, "y": 294, "flux": 28626, "ra": "260.8019", "dec": "-32.9506"}, 
+    {"x": 682, "y": 39, "flux": 28508, "ra": "260.9057", "dec": "-32.8959"}, 
+    {"x": 972, "y": 426, "flux": 27457, "ra": "260.9801", "dec": "-32.9799"}, 
+    {"x": 480, "y": 557, "flux": 26744, "ra": "260.8529", "dec": "-33.0077"}, 
+    {"x": 812, "y": 291, "flux": 25618, "ra": "260.9389", "dec": "-32.9506"}, 
+    {"x": 443, "y": 206, "flux": 25483, "ra": "260.8438", "dec": "-32.9317"}
+  ]
+}