test_io.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:
    2023-12-15, Bo Zhang, add TestFitsHeaderOps.test_append_header
    2023-12-15, Bo Zhang, add TestFitsHeaderOps.test_reformat_header
"""

import os
import unittest

from astropy.io import fits
from astropy import table

from csst_common import io


class TestFitsHeaderOps(unittest.TestCase):
    def test_append_header(self):
        h1 = fits.Header()
        h2 = fits.Header()

        h1.set("A", 1, "comment")
        h1.set("B", 1, "comment")
        h1.set("C", 1, "comment")
        h1.add_comment("=" * 72, before="A")
        h1.add_comment("one", before="A")
        h1.add_comment("=" * 72, before="A")

        h2.set("B", 2, "comment")
        h2.set("C", 2, "comment")
        h2.set("D", 2, "comment")
        h2.add_comment("=" * 72, before="B")
        h2.add_comment("another", before="B")
        h2.add_comment("=" * 72, before="B")

        h_update = io.append_header(h1, h2, duplicates="update")
        self.assertEqual(
            tuple(h_update.keys()),
            (
                "COMMENT",
                "COMMENT",
                "COMMENT",
                "A",
                "B",
                "C",
                "COMMENT",
                "COMMENT",
                "COMMENT",
                "D",
            ),
            "update mode failed",
        )
        self.assertEqual(h_update["A"], 1)
        self.assertEqual(h_update["B"], 2)
        self.assertEqual(h_update["C"], 2)
        self.assertEqual(h_update["D"], 2)

        h_delete = io.append_header(h1, h2, duplicates="delete")
        self.assertEqual(
            tuple(h_delete.keys()),
            (
                "COMMENT",
                "COMMENT",
                "COMMENT",
                "A",
                "COMMENT",
                "COMMENT",
                "COMMENT",
                "B",
                "C",
                "D",
            ),
            "delete mode failed",
        )
        self.assertEqual(h_delete["A"], 1)
        self.assertEqual(h_delete["B"], 2)
        self.assertEqual(h_delete["C"], 2)
        self.assertEqual(h_delete["D"], 2)

    def test_reformat_header(self):
        h = fits.Header()
        h.add_comment("A")
        h.add_comment("B")
        h.add_comment("B")
        h.add_comment("C")
        h.add_comment("X")
        h.set("A", 1)
        h.set("SIMPLE", True)
        h.set("NAXIS1", 1)

        h_rfmt = io.reformat_header(h, strip=True, comment="WCS info")
        self.assertEqual(
            tuple(h_rfmt.keys()),
            ("COMMENT", "COMMENT", "COMMENT", "A", "NAXIS1"),
        )

        h_rfmt = io.reformat_header(h, strip=False, comment="WCS info")
        self.assertEqual(
            tuple(h_rfmt.keys()),
            ("COMMENT", "COMMENT", "COMMENT", "A", "SIMPLE", "NAXIS1"),
        )

    def test_delete_section(self):
        h = fits.Header()
        h.append(("COMMENT", "A", ""), bottom=True)
        h.append(("COMMENT", "C", ""), bottom=True)
        h.append(("COMMENT", "B", ""), bottom=True)
        h.append(("SIMPLE", True, ""), bottom=True)
        h.append(("NAXIS1", 1, ""), bottom=True)
        h.append(("COMMENT", "=" * 72, ""), bottom=True)
        h.append(("COMMENT", "WCS", ""), bottom=True)
        h.append(("COMMENT", "=" * 72, ""), bottom=True)
        h.append(("A", 1, ""), bottom=True)
        h.append(("COMMENT", "=" * 72, ""), bottom=True)
        h.append(("COMMENT", "=" * 72, ""), bottom=True)

        h_del = io.delete_section(h, title="WCS")
        self.assertEqual(len(h_del.cards), 7)

    def test_generate_meta(self):
        h = io.generate_meta()
        self.assertIsInstance(h, dict)

    def test_append_meta(self):
        hdulist = fits.HDUList([fits.PrimaryHDU()])
        meta = io.generate_meta(obs_id="123456")
        hdulist = io.append_meta(hdulist, meta)
        self.assertIn("META", set(hdulist[0].header.keys()))
        meta = io.extract_meta(hdulist)
        self.assertEqual(meta["obs_id"], "123456")

    def test_append_meta_to_table_hdu(self):
        t = table.Table([{"a": 1, "b": 2}] * 5)
        hdu = io.table_to_hdu(t)
        meta = io.generate_meta(obs_id="123456")
        hdu = io.append_meta_to_table_hdu(hdu, meta)
        self.assertEqual(
            set(io.hdu_to_table(hdu).colnames), {"a", "b", "healpix", "data_uuid"}
        )