test_fits.py

"""
Identifier:     tests/test_fits_header_ops.py
Name:           test_fits_header_ops.py
Description:    test append_header
Author:         Bo Zhang
Created:        2023-12-15
Modified-History:
    2025-10-23, Bo Zhang, add tests for fits operations
"""

import os
import unittest

import numpy as np
from csst_common import fits

test_fits_file = os.path.join(
    os.environ["UNIT_TEST_DATA_ROOT"],
    "csst_common/test_fits/hl.fits",
)


class TestFitsHeaderOps(unittest.TestCase):

    def setUp(self):
        hl = fits.HDUList(
            [
                fits.PrimaryHDU(),
                fits.ImageHDU(data=np.random.randn(5, 5)),
            ]
        )
        hl.writeto(test_fits_file, overwrite=True)

    def test_fits_open(self):
        hl = fits.open(test_fits_file)
        self.assertIsInstance(hl, fits.HDUList)

    def test_fits_header(self):
        header = fits.getheader(test_fits_file, ext=0)
        self.assertIsInstance(header, fits.Header)

    def test_fits_getval(self):
        val = fits.getval(test_fits_file, ext=0, keyword="SIMPLE")
        self.assertEqual(val, True)

    def test_fits_getdata(self):
        data = fits.getdata(test_fits_file, ext=1)
        self.assertIsInstance(data, np.ndarray)
        self.assertEqual(data.shape, (5, 5))

    def test_fits_to_bytes(self):
        hl = fits.open(test_fits_file)
        fits_bytes = hl.to_bytes()
        self.assertIsInstance(fits_bytes, bytes)

    def test_prepare_fits_files_for_ingestion(self):
        ingestion_files = fits.prepare_fits_files_for_ingestion([test_fits_file])
        self.assertIsInstance(ingestion_files, list)
        self.assertEqual(len(ingestion_files), 1)
        self.assertEqual(set(ingestion_files[0].keys()), {"file_name", "file_content"})