noaa-goes-visualization/tests/test_fitsio.py

import numpy as np
import pytest
from astropy.io import fits

from conftest import solar_disc, write_fits
from suvi import fitsio


@pytest.fixture
def disc():
    return solar_disc()


def corrupt(path, tmp_path, name, transform):
    raw = open(path, "rb").read()
    out = str(tmp_path / name)
    with open(out, "wb") as handle:
        handle.write(transform(raw))
    return out


# ------------------------------------------------------------------ value parsing


@pytest.mark.parametrize(
    "raw,expected",
    [
        ("                   T / comment", True),
        ("                   F", False),
        ("                 171 / [angstrom]", 171),
        ("                 2.5 / arcsec", 2.5),
        ("7.40304005902818E-05 / [W m-2]", pytest.approx(7.40304005902818e-05)),
        ("1.5D3 / fortran exponent", 1500.0),
        ("'UTC     '           / principal time system", "UTC"),
        ("'W m-2 sr-1'", "W m-2 sr-1"),
        ("'0       '           / data unit checksum", "0"),
        ("'it''s quoted'       / embedded quote", "it's quoted"),
        ("'a / b'              / slash inside string", "a / b"),
        ("", None),
        ("              -999999", -999999),
    ],
)
def test_parse_value(raw, expected):
    assert fitsio._parse_value(raw) == expected


# ------------------------------------------------------------------ header reading


def test_reads_header_of_a_compressed_image(good_frame):
    cards, error = fitsio.read_header(good_frame)
    assert error is None
    assert cards["WAVELNTH"] == 171
    assert cards["ZIMAGE"] is True
    assert cards["ZNAXIS1"] == 1280
    assert cards["DEGRADED"] is False
    assert cards["IMG_MEAN"] > 0


def test_reads_header_of_an_uncompressed_image(tmp_path, disc):
    path = write_fits(str(tmp_path / "plain.fits"), disc, compress=False)
    cards, error = fitsio.read_header(path)
    assert error is None
    assert cards["NAXIS"] == 2 and cards["NAXIS1"] == 1280
    assert cards["WAVELNTH"] == 171


def test_header_metadata_projects_onto_db_columns(good_frame):
    values, error = fitsio.scan_header(good_frame)
    assert error is None
    # Booleans become 0/1 so SQL can filter on them.
    assert values["degraded"] == 0 and values["eclipse"] == 0 and values["empty"] == 0
    assert values["wavelnth"] == 171
    assert isinstance(values["img_mean"], float)
    assert isinstance(values["datasum"], str)
    assert set(values) <= set(fitsio.db.HEADER_FIELDS)


def test_header_metadata_drops_wrongly_typed_cards():
    values = fitsio.header_metadata({"IMG_MEAN": "not a number", "WAVELNTH": 171})
    assert "img_mean" not in values
    assert values["wavelnth"] == 171


def test_flags_are_read_from_a_degraded_frame(tmp_path, disc):
    """The real eclipse failure mode: near-zero radiance with DEGRADED/ECLIPSE set."""
    path = write_fits(
        str(tmp_path / "eclipse.fits"),
        disc * 1e-4,
        headers={"DEGRADED": True, "ECLIPSE": 2},
    )
    values, error = fitsio.scan_header(path)
    assert error is None
    assert values["degraded"] == 1 and values["eclipse"] == 2
    assert values["img_mean"] < 1e-3


# ----------------------------------------------------------------------- integrity


def test_empty_file(tmp_path):
    path = str(tmp_path / "empty.fits")
    open(path, "wb").close()
    cards, error = fitsio.read_header(path)
    assert cards == {} and error == "empty file"


def test_missing_file(tmp_path):
    cards, error = fitsio.read_header(str(tmp_path / "absent.fits"))
    assert cards == {} and "unreadable" in error


def test_not_a_fits_file(tmp_path):
    path = str(tmp_path / "junk.fits")
    open(path, "wb").write(b"X" * 5000)
    cards, error = fitsio.read_header(path)
    assert cards == {} and "missing SIMPLE" in error


def test_shorter_than_one_block(tmp_path, good_frame):
    path = corrupt(good_frame, tmp_path, "tiny.fits", lambda raw: raw[:100])
    _, error = fitsio.read_header(path)
    assert "truncated" in error


def test_truncation_not_on_a_block_boundary_is_reported(tmp_path, good_frame):
    path = corrupt(good_frame, tmp_path, "odd.fits", lambda raw: raw[:-100])
    _, error = fitsio.read_header(path)
    assert "not a multiple of 2880" in error


def test_header_without_end_card(tmp_path, good_frame):
    """A header whose END is gone must not be parsed as if it were complete."""
    path = corrupt(
        good_frame, tmp_path, "noend.fits", lambda raw: raw.replace(b"END     ", b"XXX     ")
    )
    cards, error = fitsio.read_header(path)
    assert cards == {} and "END" in error


def test_primary_only_file_has_no_image_hdu(tmp_path):
    path = str(tmp_path / "primary.fits")
    fits.HDUList([fits.PrimaryHDU()]).writeto(path, overwrite=True)
    cards, error = fitsio.read_header(path)
    assert "no image HDU" in error


def test_max_header_bytes_is_respected(good_frame):
    """A malformed file must not pull unbounded data into memory."""
    cards, error = fitsio.read_header(good_frame, max_bytes=fitsio.BLOCK)
    assert cards == {} or error is not None


# ------------------------------------------------------------------------ datasums


def test_datasum_matches_the_stored_keyword(good_frame):
    assert fitsio.verify_datasums(good_frame) is None


def test_datasum_detects_a_flipped_bit(tmp_path, good_frame):
    def flip(raw):
        offset = len(raw) - fitsio.BLOCK  # inside the final data unit
        return raw[:offset] + bytes([raw[offset] ^ 0xFF]) + raw[offset + 1 :]

    path = corrupt(good_frame, tmp_path, "flipped.fits", flip)
    assert "datasum mismatch" in fitsio.verify_datasums(path)


def test_datasum_detects_truncation(tmp_path, good_frame):
    path = corrupt(good_frame, tmp_path, "half.fits", lambda raw: raw[: len(raw) // 2])
    assert "truncated" in fitsio.verify_datasums(path)


def test_datasum_reports_when_no_keywords_are_present(tmp_path, disc):
    path = write_fits(str(tmp_path / "nosum.fits"), disc, checksum=False)
    assert fitsio.verify_datasums(path) == "no DATASUM keywords present"


def test_datasum_on_a_non_fits_file(tmp_path):
    path = str(tmp_path / "junk.fits")
    open(path, "wb").write(b"nope" * 1000)
    assert fitsio.verify_datasums(path) == "not a FITS file"


def test_datasum_is_the_fits_ones_complement_sum():
    # Two words that overflow 32 bits, forcing an end-around carry.
    payload = (0xFFFFFFFF).to_bytes(4, "big") + (0x00000002).to_bytes(4, "big")
    assert fitsio.datasum(payload) == 2
    assert fitsio.datasum(b"") == 0
    # A trailing partial word is zero-padded on the right: 0x00000001 + 0x01000000.
    assert fitsio.datasum(b"\x00\x00\x00\x01\x01") == 0x01000001


# --------------------------------------------------------------------- image reads


def test_reads_pixel_data(good_frame, disc):
    data, error = fitsio.read_image(good_frame)
    assert error is None
    assert data.shape == (1280, 1280) and data.dtype == np.float32
    assert np.allclose(data, disc, atol=1e-3)


def test_zblank_becomes_nan(tmp_path, disc):
    holed = disc.copy()
    holed[10:20, 10:20] = fitsio.ZBLANK
    path = write_fits(str(tmp_path / "blank.fits"), holed)
    data, error = fitsio.read_image(path)
    assert error is None
    assert np.isnan(data[10:20, 10:20]).all()
    assert np.isfinite(data[500:600, 500:600]).all()


def test_image_read_reports_a_missing_array(tmp_path):
    path = str(tmp_path / "primary.fits")
    fits.HDUList([fits.PrimaryHDU()]).writeto(path, overwrite=True)
    data, error = fitsio.read_image(path)
    assert data is None and "no 2-D image array" in error


def test_image_read_reports_corruption_instead_of_raising(tmp_path, good_frame):
    path = corrupt(good_frame, tmp_path, "broken.fits", lambda raw: raw[: len(raw) // 2])
    data, error = fitsio.read_image(path)
    assert data is None and error


def test_quiet_astropy_is_callable():
    fitsio.quiet_astropy()