Upload New File

tests/test_io.py

+import unittest
+from unittest import mock
+from CpicImgSim.io import obsid_parser, primary_hdu, frame_header, save_fits_simple
+import CpicImgSim.io as io
+from astropy.io import fits
+import numpy as np
+import yaml
+
+cstar = {
+    'magnitude': 5,
+    'ra': '0d',
+    'dec': '0d',
+    'sptype': 'G0III',
+    'distance': 10,
+    'mag_input_band': 'f661'
+}
+
+params = {
+    'target': {'cstar': cstar},
+    'skybg': 21,
+    'expt': 20,
+    'nframe': 5,
+    'band': 'f661',
+    'emgain': 10,
+    'obsid': '51012345678',
+    'rotation': 20,
+    'shift': [0, 1],
+}
+
+
+class TestIO(unittest.TestCase):
+    def test_obsid_parser(self):
+        self.assertRaises(ValueError, obsid_parser, '20190101')
+        self.assertRaises(ValueError, obsid_parser, '123456789012')
+
+        self.assertEqual(obsid_parser('50012345678'), 'BIAS')
+        self.assertEqual(obsid_parser('50112345678'), 'DARK')
+        self.assertEqual(obsid_parser('50212345678'), 'FLAT')
+        self.assertEqual(obsid_parser('50312345678'), 'BKGD')
+        self.assertEqual(obsid_parser('51012345678'), 'SCIE')
+
+    def test_primary_hdu(self):
+        hdu1 = primary_hdu(params, {}, filename_output=False)
+        hdu2, folder, filename = primary_hdu(params, {}, filename_output=True)
+        self.assertIsInstance(hdu1, fits.PrimaryHDU)
+        self.assertIsInstance(hdu2, fits.PrimaryHDU)
+        self.assertIsInstance(folder, str)
+        self.assertIsInstance(filename, str)
+
+    def test_frame_header(self):
+        header = frame_header(params, 1, '2021-01-01T00:00:00')
+        self.assertEqual(header['IMGINDEX'], 1)
+        self.assertIsInstance(header, fits.Header)
+
+    def test_write_fits(self):
+        images = np.zeros((5, 10, 10))
+        yaml_str = """
+obsid: 51012345678
+expt: 300
+nframe: 10
+band: "f661"
+shift: [0, 0]
+rotation: 0
+emgain: 100
+skybg: 21
+
+target:
+    cstar:
+        ra: "10.684792d"
+        dec: "41.26917d"
+        sptype: "M0.5"
+        magnitude: 3.4
+    planets:
+        - ra: 10.684792
+          dec: 41.26917
+          sptype: "M0.5"
+          magnitude: 3.4
+"""
+        print(io.tmp_folder_path)
+        parameters = yaml.load(yaml_str, Loader=yaml.FullLoader)
+        mock_fits_writeto = mock.MagicMock()
+        io.fits.writeto = mock_fits_writeto
+        io.tmp_folder_path = 'test_folder_for_unit_test'
+        io.save_fits_simple(images, parameters)
+        output_name = mock_fits_writeto.call_args[0][0]
+        self.assertEqual(
+            output_name[:len(io.tmp_folder_path)], io.tmp_folder_path)
+
+
+if __name__ == '__main__':
+    unittest.main()