test_io.py 4.52 KB
Newer Older
BO ZHANG's avatar
BO ZHANG committed
1
2
3
4
5
6
7
"""
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:
BO ZHANG's avatar
tweaks    
BO ZHANG committed
8
9
    2023-12-15, Bo Zhang, add TestFitsHeaderOps.test_append_header
    2023-12-15, Bo Zhang, add TestFitsHeaderOps.test_reformat_header
BO ZHANG's avatar
BO ZHANG committed
10
"""
BO ZHANG's avatar
BO ZHANG committed
11

BO ZHANG's avatar
BO ZHANG committed
12
import os
13
import unittest
BO ZHANG's avatar
BO ZHANG committed
14

15
16
from astropy.io import fits

BO ZHANG's avatar
tweaks    
BO ZHANG committed
17
18
19
20
21
22
from csst_common.io import (
    append_header,
    reformat_header,
    delete_section,
    get_proc_info,
    generate_meta,
BO ZHANG's avatar
BO ZHANG committed
23
    append_meta,
BO ZHANG's avatar
BO ZHANG committed
24
    extract_meta,
25
    fits_open,
BO ZHANG's avatar
tweaks    
BO ZHANG committed
26
)
BO ZHANG's avatar
BO ZHANG committed
27

28

BO ZHANG's avatar
tweaks    
BO ZHANG committed
29
class TestFitsHeaderOps(unittest.TestCase):
30
31
32
33
34
    def test_append_header(self):
        h1 = fits.Header()
        h2 = fits.Header()

        h1.set("A", 1, "comment")
BO ZHANG's avatar
BO ZHANG committed
35
36
        h1.set("B", 1, "comment")
        h1.set("C", 1, "comment")
37
38
39
40
        h1.add_comment("=" * 72, before="A")
        h1.add_comment("one", before="A")
        h1.add_comment("=" * 72, before="A")

BO ZHANG's avatar
BO ZHANG committed
41
42
43
        h2.set("B", 2, "comment")
        h2.set("C", 2, "comment")
        h2.set("D", 2, "comment")
44
45
46
47
        h2.add_comment("=" * 72, before="B")
        h2.add_comment("another", before="B")
        h2.add_comment("=" * 72, before="B")

BO ZHANG's avatar
BO ZHANG committed
48
        h_update = append_header(h1, h2, duplicates="update")
49
        self.assertEqual(
BO ZHANG's avatar
BO ZHANG committed
50
            tuple(h_update.keys()),
51
52
53
54
55
56
            (
                "COMMENT",
                "COMMENT",
                "COMMENT",
                "A",
                "B",
BO ZHANG's avatar
BO ZHANG committed
57
                "C",
58
59
60
                "COMMENT",
                "COMMENT",
                "COMMENT",
BO ZHANG's avatar
BO ZHANG committed
61
                "D",
62
63
64
            ),
            "update mode failed",
        )
BO ZHANG's avatar
BO ZHANG committed
65
66
67
68
69
70
        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")
71
        self.assertEqual(
BO ZHANG's avatar
BO ZHANG committed
72
            tuple(h_delete.keys()),
73
74
75
76
77
78
79
80
81
82
            (
                "COMMENT",
                "COMMENT",
                "COMMENT",
                "A",
                "COMMENT",
                "COMMENT",
                "COMMENT",
                "B",
                "C",
BO ZHANG's avatar
BO ZHANG committed
83
                "D",
84
85
86
            ),
            "delete mode failed",
        )
BO ZHANG's avatar
BO ZHANG committed
87
88
89
90
        self.assertEqual(h_delete["A"], 1)
        self.assertEqual(h_delete["B"], 2)
        self.assertEqual(h_delete["C"], 2)
        self.assertEqual(h_delete["D"], 2)
BO ZHANG's avatar
BO ZHANG committed
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

    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"),
        )
BO ZHANG's avatar
BO ZHANG committed
114
115
116

    def test_delete_section(self):
        h = fits.Header()
BO ZHANG's avatar
BO ZHANG committed
117
118
119
120
121
122
123
124
125
126
127
        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)
BO ZHANG's avatar
BO ZHANG committed
128
129
130

        h_del = delete_section(h, title="WCS")
        self.assertEqual(len(h_del.cards), 7)
BO ZHANG's avatar
BO ZHANG committed
131
132
133

    def test_get_proc_info(self):
        proc_info = get_proc_info()
BO ZHANG's avatar
tweaks    
BO ZHANG committed
134
        self.assertEquals(proc_info["BUILD"], "__unset__")
BO ZHANG's avatar
tweaks    
BO ZHANG committed
135
136
137
138

    def test_generate_meta(self):
        h = generate_meta()
        self.assertIsInstance(h, fits.Header)
BO ZHANG's avatar
tweaks    
BO ZHANG committed
139

BO ZHANG's avatar
BO ZHANG committed
140
    def test_append_meta(self):
BO ZHANG's avatar
BO ZHANG committed
141
        hdulist = fits.HDUList([fits.PrimaryHDU()])
BO ZHANG's avatar
BO ZHANG committed
142
        hdulist = append_meta(hdulist, obs_id="123456")
BO ZHANG's avatar
BO ZHANG committed
143
144
145
        self.assertIn("META", set(hdulist[0].header.keys()))
        meta = extract_meta(hdulist)
        self.assertEqual(meta["obs_id"], "123456")
146
147

    def test_fits_open(self):
BO ZHANG's avatar
BO ZHANG committed
148
149
150
151
152
153
        hl = fits_open(
            os.path.join(
                os.environ["UNIT_TEST_DATA_ROOT"],
                "csst_common/test_fits/hl.fits",
            )
        )
154
        self.assertTrue(isinstance(hl, fits.HDUList))