noaa-goes-visualization/tests/conftest.py

115 lines
4.1 KiB
Python

"""Shared fixtures.
Tests build their own FITS files rather than reading the archive, so the suite runs
anywhere and cannot be perturbed by (or perturb) the real data.
"""
import os
import sys
import numpy as np
import pytest
from astropy.io import fits
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from suvi import paths # noqa: E402
def solar_disc(size=1280, radius=386, centre=None, peak=1.0, background=0.01,
active_region=True):
"""A synthetic SUVI-like frame: a bright limb-darkened disc on a dim corona.
An off-centre active region is included by default, because a perfectly
symmetric disc is invariant under rotation and reflection and would silently
pass tests for corruptions that flip or rotate the frame. Real SUVI frames are
never symmetric, and that asymmetry is what displaces the brightness centroid.
"""
centre = centre if centre is not None else ((size - 1) / 2.0, (size - 1) / 2.0)
yy, xx = np.mgrid[0:size, 0:size].astype(np.float32)
r = np.hypot(xx - centre[0], yy - centre[1])
disc = np.zeros((size, size), dtype=np.float32)
inside = r < radius
# Limb darkening keeps the radial profile non-trivial so geometry checks bite.
disc[inside] = peak * (0.4 + 0.6 * np.sqrt(np.clip(1 - (r[inside] / radius) ** 2, 0, 1)))
corona = background * np.exp(-np.clip(r - radius, 0, None) / (radius / 4.0))
image = disc + corona
if active_region:
spot_r = np.hypot(xx - (centre[0] + radius * 0.35), yy - (centre[1] - radius * 0.3))
image = image + peak * 0.25 * np.exp(-((spot_r / (radius * 0.16)) ** 2)) * inside
return image.astype(np.float32)
def write_fits(path, data, headers=None, compress=True, checksum=True):
"""Write a SUVI-shaped FITS file: empty primary + (compressed) image extension."""
header = fits.Header()
header["DATE-BEG"] = "2024-05-10T00:00:00.000"
header["DATE-OBS"] = "2024-05-10T00:00:21.000"
header["DATE-END"] = "2024-05-10T00:00:31.000"
header["WAVELNTH"] = 171
header["EXPTIME"] = 1.0
header["CRPIX1"] = (data.shape[1] + 1) / 2.0
header["CRPIX2"] = (data.shape[0] + 1) / 2.0
header["CDELT1"] = 2.5
header["CDELT2"] = 2.5
header["CROTA"] = 0.0
header["DSUN_OBS"] = 148781338180.972
header["SOLAR_B0"] = -7.170855
header["DIAM_SUN"] = 771.9772
header["YAW_FLIP"] = 0
header["ECLIPSE"] = 0
header["EMPTY"] = False
header["DEGRADED"] = False
header["NUM_IMGS"] = 2
header["N_LONG"] = 1
header["N_SHORT"] = 0
header["N_SH_FL"] = 1
finite = data[np.isfinite(data)]
header["IMG_MIN"] = float(finite.min()) if finite.size else 0.0
header["IMG_MAX"] = float(finite.max()) if finite.size else 0.0
header["IMG_MEAN"] = float(finite.mean()) if finite.size else 0.0
header["IMG_SDEV"] = float(finite.std()) if finite.size else 0.0
header["BUNIT"] = "W m-2 sr-1"
for key, value in (headers or {}).items():
header[key] = value
if compress:
hdu = fits.CompImageHDU(data=data, header=header, compression_type="GZIP_2")
else:
hdu = fits.ImageHDU(data=data, header=header)
hdulist = fits.HDUList([fits.PrimaryHDU(), hdu])
os.makedirs(os.path.dirname(path), exist_ok=True)
hdulist.writeto(path, overwrite=True, checksum=checksum)
return path
@pytest.fixture
def archive(tmp_path):
"""An empty archive root with SUVI_DATA_ROOT pointed at it."""
root = tmp_path / "Data"
root.mkdir()
old = os.environ.get("SUVI_DATA_ROOT")
os.environ["SUVI_DATA_ROOT"] = str(root)
yield root
if old is None:
os.environ.pop("SUVI_DATA_ROOT", None)
else:
os.environ["SUVI_DATA_ROOT"] = old
@pytest.fixture
def db_path(tmp_path):
return str(tmp_path / "index.sqlite")
@pytest.fixture
def frame_name():
return paths.parse_frame_filename(
"dr_suvi-l2-ci171_g16_s20240510T000000Z_e20240510T000400Z_v1-0-2.fits"
)
@pytest.fixture
def good_frame(tmp_path):
"""A healthy synthetic frame on disk."""
return write_fits(str(tmp_path / "frames" / "good.fits"), solar_disc())