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.io import (
    append_header,
    reformat_header,
    delete_section,
    get_proc_info,
    generate_meta,
    append_meta,
    extract_meta,
    add_meta_to_table,
)


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 = 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 = 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 = reformat_header(h, strip=True, comment="WCS info")
        self.assertEqual(
            tuple(h_rfmt.keys()),
            ("COMMENT", "COMMENT", "COMMENT", "A", "NAXIS1"),
        )

        h_rfmt = 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 = delete_section(h, title="WCS")
        self.assertEqual(len(h_del.cards), 7)

    def test_get_proc_info(self):
        proc_info = get_proc_info()
        self.assertEquals(proc_info["BUILD"], "")

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

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

    def test_add_meta_to_table(self):
        t = table.Table([{"a": 1, "b": 2}] * 5)
        meta = generate_meta(obs_id="123456")
        t = add_meta_to_table(t, meta)
        self.assertEqual(set(t.colnames), {"a", "b", "healpix", "data_uuid"})