Commit 9d85a905 authored by Chen Yili's avatar Chen Yili
Browse files

Upload New File

parent ceb06b5f
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()
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment