test_io.py

import unittest
from unittest import mock

from csst_cpic_sim.config import config

from csst_cpic_sim.io import obsid_parser, primary_hdu, frame_header, set_up_logger
from csst_cpic_sim.camera import CpicVisEmccd
import csst_cpic_sim.io as io
from astropy.io import fits
import numpy as np
import yaml
import logging

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,
    'emset': 10,
    'obsid': '40112345678',
    'rotation': 20,
    'shift': [0, 1],
    'EXPSTART': '2020-01-01T00:00:00.000',
    'EXPEND': '2020-01-02T00:00:00.000',
    'frame_info': [
        {
            'expt_start': 0,
            'expt_end': 20,
            'platescale': 1.0,
            'iwa': 0,
            'chiptemp': 1}
    ]
}

camera = CpicVisEmccd()


class TestIO(unittest.TestCase):
    def test_obsid_parser(self):
        self.assertRaises(ValueError, obsid_parser, '20190101')
        self.assertRaises(ValueError, obsid_parser, '123456789012')
        self.assertRaises(ValueError, obsid_parser, '51012345678')

        self.assertEqual(obsid_parser('42012345678'), 'BIAS')
        self.assertEqual(obsid_parser('42112345678'), 'DARK')
        self.assertEqual(obsid_parser('42212345678'), 'FLAT')
        self.assertEqual(obsid_parser('42312345678'), 'BKG')
        self.assertEqual(obsid_parser('42412345678'), 'LASER')
        self.assertEqual(obsid_parser('40112345678'), 'SCI')
        self.assertEqual(obsid_parser('40212345678'), 'DSF')
        self.assertEqual(obsid_parser('41012345678'), 'CALS')
        self.assertEqual(obsid_parser('40312345678'), 'DEFT')

    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):
        hdu1 = primary_hdu(params, {}, filename_output=False)
        header = frame_header(params, 0, hdu1.header, camera.__dict__)
        self.assertEqual(header['IMGINDEX'], 1)
        self.assertIsInstance(header, fits.Header)

    @mock.patch("os.makedirs")
    @mock.patch("astropy.io.fits.writeto")
    def test_write_fits(self, patch_fits, patch_mkdir):

        log = set_up_logger(config['log_dir'])
        self.assertIsInstance(log, logging.Logger)

        self.assertRaises(FileNotFoundError, set_up_logger, 'new folder')

        patch_mkdir.assert_called_once_with('new folder')

        images = np.zeros((5, 10, 10))
        yaml_str = """
obsid: 42012345678
expt: 300
nframe: 10
band: "f661"
shift: [0, 0]
rotation: 0
emgain: 100
emset: -1
skybg: 21
expstart:

target:
    cstar:
        ra: "10.684792d"
        dec: "41.26917d"
        sptype: "M0.5"
        magnitude: 3.4
    objects:
        - ra: 10.684792
          dec: 41.26917
          sptype: "M0.5"
          magnitude: 3.4
"""
        parameters = yaml.load(yaml_str, Loader=yaml.FullLoader)
        tmp_folder_path = 'test_folder_for_unit_test'
        io.save_fits_simple(images, parameters, output_folder=tmp_folder_path)
        # patch_mkdir.assert_called_twice()

        dirname = patch_mkdir.call_args[0][0]
        self.assertEqual(dirname, tmp_folder_path)
        # patch_mkdir.assert_called_once_with(tmp_folder_path)
        output_name = patch_fits.call_args[0][0]
        self.assertEqual(
            output_name[:len(tmp_folder_path)], tmp_folder_path)


if __name__ == '__main__':
    unittest.main()