339 lines
11 KiB
Go
339 lines
11 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package io
|
||
|
||
import (
|
||
"encoding/binary"
|
||
"math"
|
||
"os"
|
||
"path/filepath"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
"strings"
|
||
"testing"
|
||
)
|
||
|
||
func fitsTempPath(t *testing.T) string {
|
||
t.Helper()
|
||
return filepath.Join(t.TempDir(), "image.fits")
|
||
}
|
||
|
||
// padCard right-pads a card body with spaces to the full 80 bytes.
|
||
func padCard(text string) []byte {
|
||
return []byte(text + strings.Repeat(" ", 80-len(text)))
|
||
}
|
||
|
||
// TestFITSRoundTripFloat64 moves a rank-2 float64 image with negative
|
||
// and non-round values through the file format and back, values and
|
||
// header strings included.
|
||
func TestFITSRoundTripFloat64(t *testing.T) {
|
||
a := mustFloats(t, []float64{
|
||
1.5, -2.25, 3.125, 4,
|
||
-5.5, 6.75, -7.875, 8,
|
||
9.25, -10.5, 11.125, -12,
|
||
}, 3, 4)
|
||
path := fitsTempPath(t)
|
||
headers := map[string]string{
|
||
"OBJECT": "M31",
|
||
"OBSERVER": "petr's dome",
|
||
"EXPTIME": "600",
|
||
}
|
||
if err := SaveFITS(path, a, headers); err != nil {
|
||
t.Fatalf("SaveFITS: %v", err)
|
||
}
|
||
back, hdr, err := LoadFITS(path)
|
||
if err != nil {
|
||
t.Fatalf("LoadFITS: %v", err)
|
||
}
|
||
if back.Dtype() != core.Float || back.NDim() != 2 || back.Shape()[0] != 3 || back.Shape()[1] != 4 {
|
||
t.Fatalf("shape/dtype mismatch: %v %s", back.Shape(), back.Dtype())
|
||
}
|
||
for i := range a.Len() {
|
||
if back.FloatAt(i) != a.FloatAt(i) {
|
||
t.Fatalf("value[%d] = %g, want %g", i, back.FloatAt(i), a.FloatAt(i))
|
||
}
|
||
}
|
||
for key, want := range map[string]string{
|
||
"OBJECT": "M31",
|
||
"OBSERVER": "petr's dome",
|
||
"EXPTIME": "600",
|
||
} {
|
||
if hdr[key] != want {
|
||
t.Fatalf("header %q = %q, want %q", key, hdr[key], want)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestFITSAxisConvention pins the wire format against the raw bytes:
|
||
// NAXIS1 must carry the fastest (last Go) axis and the payload must
|
||
// be big-endian in flat row-major order.
|
||
func TestFITSAxisConvention(t *testing.T) {
|
||
a := mustFloats(t, []float64{1, 2, 3, 4, 5, 6}, 2, 3)
|
||
path := fitsTempPath(t)
|
||
if err := SaveFITS(path, a, nil); err != nil {
|
||
t.Fatalf("SaveFITS: %v", err)
|
||
}
|
||
raw, err := os.ReadFile(path)
|
||
if err != nil {
|
||
t.Fatalf("ReadFile: %v", err)
|
||
}
|
||
if len(raw)%2880 != 0 {
|
||
t.Fatalf("file length %d is not a multiple of 2880", len(raw))
|
||
}
|
||
header := raw[:2880]
|
||
for _, want := range []string{
|
||
"SIMPLE = T",
|
||
"BITPIX = -64",
|
||
"NAXIS = 2",
|
||
"NAXIS1 = 3",
|
||
"NAXIS2 = 2",
|
||
} {
|
||
if !strings.Contains(string(header), want) {
|
||
t.Fatalf("header misses %q", want)
|
||
}
|
||
}
|
||
endAt := strings.Index(string(header), "END")
|
||
if endAt < 0 {
|
||
t.Fatal("header has no END card")
|
||
}
|
||
var word [8]byte
|
||
for i := range 6 {
|
||
binary.BigEndian.PutUint64(word[:], math.Float64bits(float64(i+1)))
|
||
got := raw[2880+i*8 : 2880+i*8+8]
|
||
if string(got) != string(word[:]) {
|
||
t.Fatalf("payload word %d = % x, want % x", i, got, word)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestFITSRoundTripFloat32 keeps the float32 element type through the
|
||
// round trip, which is what BITPIX −32 stores.
|
||
func TestFITSRoundTripFloat32(t *testing.T) {
|
||
a := core.New(core.Float32, 5)
|
||
for i := range 5 {
|
||
a.RawFloat32s()[i] = float32(i) * 1.25
|
||
}
|
||
path := fitsTempPath(t)
|
||
if err := SaveFITS(path, a, nil); err != nil {
|
||
t.Fatalf("SaveFITS: %v", err)
|
||
}
|
||
back, _, err := LoadFITS(path)
|
||
if err != nil {
|
||
t.Fatalf("LoadFITS: %v", err)
|
||
}
|
||
if back.Dtype() != core.Float32 {
|
||
t.Fatalf("dtype = %s, want core.Float32", back.Dtype())
|
||
}
|
||
for i := range 5 {
|
||
if back.RawFloat32s()[i] != a.RawFloat32s()[i] {
|
||
t.Fatalf("value[%d] = %g, want %g", i, back.RawFloat32s()[i], a.RawFloat32s()[i])
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestFITSSkipsValuelessCards checks COMMENT and HISTORY cards are
|
||
// tolerated and skipped rather than parsed as values.
|
||
func TestFITSSkipsValuelessCards(t *testing.T) {
|
||
a := mustFloats(t, []float64{1, 2, 3, 4}, 2, 2)
|
||
path := fitsTempPath(t)
|
||
if err := SaveFITS(path, a, map[string]string{"OBJECT": "TEST"}); err != nil {
|
||
t.Fatalf("SaveFITS: %v", err)
|
||
}
|
||
raw, err := os.ReadFile(path)
|
||
if err != nil {
|
||
t.Fatalf("ReadFile: %v", err)
|
||
}
|
||
// Splice two valueless cards in before the END card and re-pad the
|
||
// header to the block boundary. The END card is matched in full
|
||
// ("END" plus its padding) because EXTEND contains the same three
|
||
// letters.
|
||
endCard := strings.Index(string(raw), "END"+strings.Repeat(" ", 77))
|
||
if endCard < 0 {
|
||
t.Fatal("no END card")
|
||
}
|
||
spliced := append([]byte{}, raw[:endCard]...)
|
||
spliced = append(spliced, padCard("COMMENT a note without a value")...)
|
||
spliced = append(spliced, padCard("HISTORY an audit trail entry")...)
|
||
spliced = append(spliced, padCard("END")...)
|
||
for len(spliced)%2880 != 0 {
|
||
spliced = append(spliced, padCard("")...)
|
||
}
|
||
spliced = append(spliced, raw[2880:]...)
|
||
if err := os.WriteFile(path, spliced, 0o644); err != nil {
|
||
t.Fatalf("WriteFile: %v", err)
|
||
}
|
||
back, hdr, err := LoadFITS(path)
|
||
if err != nil {
|
||
t.Fatalf("LoadFITS: %v", err)
|
||
}
|
||
if back.Len() != 4 || back.FloatAt(3) != 4 {
|
||
t.Fatalf("payload damaged: %v %v", back.Shape(), back.RawFloats())
|
||
}
|
||
if hdr["OBJECT"] != "TEST" {
|
||
t.Fatalf("OBJECT = %q, want %q", hdr["OBJECT"], "TEST")
|
||
}
|
||
if _, ok := hdr["COMMENT"]; ok {
|
||
t.Fatal("COMMENT card must not enter the header map")
|
||
}
|
||
}
|
||
|
||
// TestFITSErrors covers every refusal path in the format layer.
|
||
func TestFITSErrors(t *testing.T) {
|
||
path := fitsTempPath(t)
|
||
|
||
ints, ierr := core.FromInts([]int64{1, 2, 3, 4}, 2, 2)
|
||
if ierr != nil {
|
||
t.Fatalf("FromInts: %v", ierr)
|
||
}
|
||
if err := SaveFITS(path, ints, nil); err == nil {
|
||
t.Fatal("int64 input: want an error")
|
||
}
|
||
if err := SaveFITS(path, mustComplexes(t, []complex128{1, 2, 3, 4}, 2, 2), nil); err == nil {
|
||
t.Fatal("complex input: want an error")
|
||
}
|
||
ok := mustFloats(t, []float64{1, 2}, 2)
|
||
if err := SaveFITS(path, ok, map[string]string{"TOOLONGKEYWORD": "x"}); err == nil {
|
||
t.Fatal("long keyword: want an error")
|
||
}
|
||
if err := SaveFITS(path, ok, map[string]string{"BITPIX": "x"}); err == nil {
|
||
t.Fatal("reserved keyword: want an error")
|
||
}
|
||
if err := SaveFITS(path, ok, map[string]string{"naxis2": "x"}); err == nil {
|
||
t.Fatal("NAXISn keyword: want an error")
|
||
}
|
||
if err := SaveFITS(path, ok, map[string]string{"BAD KEY": "x"}); err == nil {
|
||
t.Fatal("space in keyword: want an error")
|
||
}
|
||
if err := SaveFITS(path, ok, map[string]string{"NOTE": strings.Repeat("x", 69)}); err == nil {
|
||
t.Fatal("overlong value: want an error")
|
||
}
|
||
|
||
// A valid file to damage in every way the parser must catch.
|
||
if err := SaveFITS(path, ok, nil); err != nil {
|
||
t.Fatalf("SaveFITS: %v", err)
|
||
}
|
||
raw, err := os.ReadFile(path)
|
||
if err != nil {
|
||
t.Fatalf("ReadFile: %v", err)
|
||
}
|
||
with := func(mutate func([]byte) []byte) {
|
||
t.Helper()
|
||
if err := os.WriteFile(path, mutate(append([]byte{}, raw...)), 0o644); err != nil {
|
||
t.Fatalf("WriteFile: %v", err)
|
||
}
|
||
if _, _, err := LoadFITS(path); err == nil {
|
||
t.Fatal("damaged file: want an error")
|
||
}
|
||
}
|
||
with(func(b []byte) []byte { return b[:100] }) // no END card
|
||
with(func(b []byte) []byte { return b[:2880] }) // data truncated
|
||
with(func(b []byte) []byte { b[30] = 'F'; return b }) // SIMPLE = F
|
||
with(func(b []byte) []byte { copy(b[11:20], "XTENSION"); return b }) // wrong first card
|
||
|
||
// BITPIX 8 (unsigned bytes) is outside the supported image types.
|
||
unsupported := append(append(append(append([]byte{},
|
||
padCard("SIMPLE = T")...),
|
||
padCard("BITPIX = 8")...),
|
||
padCard("NAXIS = 1")...),
|
||
padCard("NAXIS1 = 4")...)
|
||
unsupported = append(unsupported, padCard("END")...)
|
||
unsupported = append(unsupported, make([]byte, 2880-5*80)...)
|
||
unsupported = append(unsupported, make([]byte, 2880)...)
|
||
if err := os.WriteFile(path, unsupported, 0o644); err != nil {
|
||
t.Fatalf("WriteFile: %v", err)
|
||
}
|
||
if _, _, err := LoadFITS(path); err == nil {
|
||
t.Fatal("BITPIX 8: want an error")
|
||
}
|
||
|
||
// A hostile header whose axis product overflows int must be
|
||
// refused as truncated data, never panic the process.
|
||
hostile := append(append(append(append([]byte{},
|
||
padCard("SIMPLE = T")...),
|
||
padCard("BITPIX = -64")...),
|
||
padCard("NAXIS = 2")...),
|
||
padCard("NAXIS1 = 1099511627776")...)
|
||
hostile = append(hostile, padCard("NAXIS2 = 1099511627776")...)
|
||
hostile = append(hostile, padCard("END")...)
|
||
hostile = append(hostile, make([]byte, 2880-6*80)...)
|
||
hostile = append(hostile, make([]byte, 2880)...)
|
||
if err := os.WriteFile(path, hostile, 0o644); err != nil {
|
||
t.Fatalf("WriteFile: %v", err)
|
||
}
|
||
if _, _, err := LoadFITS(path); err == nil {
|
||
t.Fatal("overflowing axis product: want an error")
|
||
}
|
||
|
||
// An XTENSION-first file is an extension, not a primary image.
|
||
ext := append([]byte{}, padCard("XTENSION= 'IMAGE '")...)
|
||
ext = append(ext, raw[80:]...)
|
||
if err := os.WriteFile(path, ext, 0o644); err != nil {
|
||
t.Fatalf("WriteFile: %v", err)
|
||
}
|
||
if _, _, err := LoadFITS(path); err == nil {
|
||
t.Fatal("extension header: want an error")
|
||
}
|
||
|
||
// The intact file still loads after all the damage around it.
|
||
if err := os.WriteFile(path, raw, 0o644); err != nil {
|
||
t.Fatalf("WriteFile: %v", err)
|
||
}
|
||
if _, _, err := LoadFITS(path); err != nil {
|
||
t.Fatalf("intact file: %v", err)
|
||
}
|
||
}
|
||
|
||
// TestCommentaryCardsAreNotValueCards pins the commentary rule: a
|
||
// COMMENT or HISTORY card may legitimately carry an "= " sequence in
|
||
// columns 9-10, and such a card must be skipped as commentary, not
|
||
// parsed as a keyword with a value.
|
||
func TestCommentaryCardsAreNotValueCards(t *testing.T) {
|
||
cards := []string{
|
||
fitsBoolCard("SIMPLE", true),
|
||
fitsIntCard("BITPIX", -64),
|
||
fitsIntCard("NAXIS", 1),
|
||
fitsIntCard("NAXIS1", 2),
|
||
fitsPadCard("COMMENT = this looks like a value card"),
|
||
fitsPadCard("HISTORY = so does this one"),
|
||
fitsStringCardRaw("OBSERVER", "tester"),
|
||
fitsEndCard(),
|
||
}
|
||
data := fitsAppendCards(nil, cards)
|
||
payload := []byte{0x3f, 0xf0, 0, 0, 0, 0, 0, 0, 0x40, 0, 0, 0, 0, 0, 0, 0} // 1.0, 2.0
|
||
data = append(data, payload...)
|
||
data = fitsAppendZeroPad(data)
|
||
|
||
path := filepath.Join(t.TempDir(), "commentary.fits")
|
||
if err := osWriteFile(path, data); err != nil {
|
||
t.Fatalf("os.WriteFile: %v", err)
|
||
}
|
||
img, headers, err := LoadFITS(path)
|
||
if err != nil {
|
||
t.Fatalf("LoadFITS: %v", err)
|
||
}
|
||
if img.Len() != 2 || img.FloatAt(0) != 1 || img.FloatAt(1) != 2 {
|
||
t.Fatalf("image = %s, want [1, 2]", img)
|
||
}
|
||
if _, ok := headers["COMMENT"]; ok {
|
||
t.Error("COMMENT was parsed as a value keyword")
|
||
}
|
||
if _, ok := headers["HISTORY"]; ok {
|
||
t.Error("HISTORY was parsed as a value keyword")
|
||
}
|
||
if headers["OBSERVER"] != "tester" {
|
||
t.Errorf("OBSERVER = %q, want tester", headers["OBSERVER"])
|
||
}
|
||
}
|
||
|
||
// TestFitsCheckKeywordRefusesCommentary pins that COMMENT and HISTORY
|
||
// are refused as user keywords: they carry no value in the format.
|
||
func TestFitsCheckKeywordRefusesCommentary(t *testing.T) {
|
||
dir := t.TempDir()
|
||
img := mustFloats(t, []float64{1, 2}, 2)
|
||
for _, kw := range []string{"COMMENT", "HISTORY"} {
|
||
if err := SaveFITS(filepath.Join(dir, strings.ToLower(kw)+".fits"), img, map[string]string{kw: "x"}); err == nil {
|
||
t.Errorf("SaveFITS accepted the reserved keyword %q", kw)
|
||
}
|
||
}
|
||
}
|