test_main.py

import unittest
import shutil
import os
# from CpicImgSim import observation_simulation, quick_run
from csst_cpic_sim.main import main, quick_run_v2, observation_simulation_from_config
from csst_cpic_sim.config import config

test_dir = os.path.dirname(__file__)
cases = os.path.join(test_dir, 'testcases')
output = os.path.join(test_dir, 'test_output')
config['output'] = output
# from unittest.mock import patch

a = []


def clear_output():
    if os.path.exists(output):
        shutil.rmtree(output)
    os.mkdir(output)


class TestMain(unittest.TestCase):
    # def test_help(self):
    #     main(argv = None)

    def test_quick_run_cmdline(self):
        clear_output()
        main([
            'quickrun',
            'demo_0_20',
            '0', '700', '1',
            '-r', '30',
            '-s', '19',
            '-ec',
            '-o', output
        ])
        file = os.listdir(output)
        self.assertEqual(len(file), 1)
        self.assertEqual(file[0][:9], 'demo_0_20')

    def test_quick_run_func(self):
        clear_output()
        quick_run_v2(
            '', 'f661', 0, 20, 1, skybg=None, camera_effect=False, output=output
        )
        file = os.listdir(output)
        self.assertEqual(len(file), 1)
        self.assertEqual(file[0][:5], 'blank')

    def test_full_run_func(self):
        clear_output()
        with open(cases+'/run_config.yaml', 'r') as fid:
            text = fid.read()

        text = text.replace('<output>', output)
        config_text = text.replace('<msc_model>', "True")
        with open(output+'/run_config_new.yaml', 'w') as fid:
            fid.write(config_text)

        observation_simulation_from_config(
            os.path.join(cases, '05_sci.yaml'),
            os.path.join(output, 'run_config_new.yaml')
        )

        file = os.listdir(output)
        self.assertEqual(len(file), 4)

        clear_output()

        observation_simulation_from_config(
            os.path.join(cases, '05_sci.yaml'),
            None
        )
        file = os.listdir(output)
        self.assertEqual(len(file), 1)
        self.assertEqual(file[0][:5], 'SCI')

        clear_output()

        config_text = text.replace('<msc_model>', "False")
        with open(output+'/run_config_new.yaml', 'w') as fid:
            fid.write(config_text)

        observation_simulation_from_config(
            os.path.join(cases, '05_sci.yaml'),
            os.path.join(output, 'run_config_new.yaml')
        )

        file = os.listdir(output)
        self.assertEqual(len(file), 4)

    def test_full_run_cmdline(self):
        clear_output()
        main(['run', os.path.join(cases, '05_sci.yaml')])
        file = os.listdir(output)
        self.assertEqual(len(file), 1)
        self.assertEqual(file[0][:5], 'SCI')

    # def test_main(self):
    #     target_example = {
    #         'name': 'test',
    #         'cstar': {
    #             'magnitude': 0,
    #             'ra': '120d',
    #             'dec': '40d',
    #             'distance': 10,
    #             'sptype': 'F0III',
    #         },
    #     }

    #     quick_run('', 10, 'f661', 1, 1, 30)
    #     quick_run('*2.4/10(3,5)/15(-4,2)', 21, 'f661', 1, 1, 30)

    #     # normal target
    #     observation_simulation(
    #         target=target_example,
    #         skybg=21,
    #         expt=1,
    #         nframe=1,
    #         band='f661',
    #         emgain=30,
    #         obsid=51012345678,
    #     )

    #     # bias-gain
    #     observation_simulation(
    #         target={},
    #         skybg=None,
    #         expt=1,
    #         nframe=1,
    #         band='f661',
    #         emgain=30,
    #         obsid=50112345678,
    #     )


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