test_fits_header_ops.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 unittest

from astropy.io import fits

from csst_common.io import append_header, reformat_header


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

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

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

        self.assertEqual(
            tuple(append_header(h1, h2, duplicates="update").keys()),
            (
                "COMMENT",
                "COMMENT",
                "COMMENT",
                "A",
                "B",
                "COMMENT",
                "COMMENT",
                "COMMENT",
                "C",
            ),
            "update mode failed",
        )
        self.assertEqual(
            tuple(append_header(h1, h2, duplicates="delete").keys()),
            (
                "COMMENT",
                "COMMENT",
                "COMMENT",
                "A",
                "COMMENT",
                "COMMENT",
                "COMMENT",
                "B",
                "C",
            ),
            "delete mode failed",
        )

    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"),
        )