noaa-goes-visualization/tests/test_render_real.py

"""The real-data renderer: judged-bad and missing bands are filled, nothing else.

No model or compositor runs here -- both are stubbed -- because what these tests
pin is the policy: which slots count as bad, which bands the model supplies,
which timestamps stay black, and that a restart repeats no work.
"""

import numpy as np
import pytest

import render_real
from conftest import solar_disc, write_fits
from suvi import cases, detectors, paths

DAY = render_real.DAY
DAY0 = 1717200000 - 1717200000 % DAY          # 2024-06-01 00:00 UTC
BANDS = tuple(paths.WAVELENGTHS)
#: Brightness per band that header_v1 accepts for a peak-1.0 synthetic disc.
BAND_SCALE = {94: 0.5, 131: 0.5, 171: 1.0, 195: 2.0, 284: 2.0, 304: 3.0}
#: Side of the stubbed frames.  render_day never looks at pixel content, and the
#: harness keeps a whole day of jobs: at full size (1280) that is ~23 GB, enough to
#: run the machine out of memory.
STUB = 8


def write_slot(tmp_path, satellite, when, dark=False):
    """Six real FITS files for one slot; `dark` gives a disc-less frame."""
    rows = []
    for wavelength in BANDS:
        name = f"g{satellite}_{wavelength}_{when}.fits"
        image = (np.zeros((1280, 1280), dtype=np.float32) if dark
                 else solar_disc(size=1280, radius=386, peak=1.0) * BAND_SCALE[wavelength])
        write_fits(str(tmp_path / name), image, headers={"WAVELNTH": wavelength})
        rows.append((satellite, wavelength, when, name, 1_500_000))
    return rows


# ----------------------------------------------------------------- windowing


def test_a_day_is_seen_through_two_days_either_side():
    t0, t1 = render_real.window_for(DAY0)
    assert t0 == DAY0 - 2 * DAY and t1 == DAY0 + 3 * DAY


def test_the_window_covers_the_stack_and_the_dual_gap_policy():
    """+/-128 slots is ~8.5 h and the policy is 24 h: a dual gap overlapping the
    day must be visible at its full length up to the limit, or a long outage
    would be judged short and invented rather than left black."""
    from suvi import availability
    t0, t1 = render_real.window_for(DAY0)
    policy = availability.MAX_DUAL_GAP * paths.CADENCE
    assert DAY0 - t0 >= policy and t1 - (DAY0 + DAY) >= policy
    assert DAY0 - t0 >= 128 * paths.CADENCE


def test_day_starts_align_to_utc_midnight_and_exclude_the_end():
    days = render_real.day_starts(DAY0 + 3600, DAY0 + 3 * DAY)
    assert days == [DAY0, DAY0 + DAY, DAY0 + 2 * DAY]


# --------------------------------------------------------- what counts as bad


def test_stub_downloads_are_treated_as_missing(tmp_path):
    rows = [(16, 171, DAY0, "a.fits", 1_500_000, "good"),
            (16, 171, DAY0 + 240, "b.fits", 5_000, None)]
    archive = render_real.archive_map(rows, str(tmp_path), BANDS)
    assert (16, 171, DAY0) in archive
    assert (16, 171, DAY0 + 240) not in archive


def test_bad_comes_from_the_stored_verdicts_and_unjudged_is_reported():
    rows = [(16, 171, DAY0, "a", 1_500_000, "good"),
            (16, 171, DAY0 + 240, "b", 1_500_000, "bad"),
            (16, 171, DAY0 + 480, "c", 1_500_000, None),
            (16, 171, DAY0 + 720, "d", 5_000, None)]          # a stub needs no verdict
    archive = render_real.archive_map(rows, "/r", BANDS)
    bad, unjudged = render_real.judged_bad(rows, archive)
    assert bad == {(16, 171, DAY0 + 240)}
    assert unjudged == {(16, 171, DAY0 + 480)}


def test_frame_rows_carry_only_the_named_runs_verdict(tmp_path):
    from suvi import db

    conn = db.connect(str(tmp_path / "index.sqlite"))
    name = paths.FrameName(satellite=16, wavelength=171, t_start=DAY0, t_end=DAY0 + 240,
                           version="1-0-2")
    frame_id = db.upsert_frame(conn, name, name.relpath(), size_bytes=1_500_000)
    mine = db.create_detector_run(conn, "judge", {"a": 1})
    other = db.create_detector_run(conn, "judge", {"a": 2})
    db.record_detections(conn, other, [(frame_id, "bad", None, None, 0)])
    conn.commit()
    rows = render_real.frame_rows(conn, DAY0, DAY0 + DAY, (16,), mine)
    assert rows == [(16, 171, DAY0, name.relpath(), 1_500_000, None)]
    db.record_detections(conn, mine, [(frame_id, "good", None, None, 0)])
    assert render_real.frame_rows(conn, DAY0, DAY0 + DAY, (16,), mine)[0][-1] == "good"


def test_bands_to_fill_covers_bad_and_absent_bands():
    archive = {(16, w, DAY0): "/x" for w in BANDS if w != 304}
    overlay = cases.Overlay(archive=archive)
    bad = {(16, 171, DAY0)}
    missing = render_real.bands_to_fill(overlay, bad, 16, DAY0, BANDS)
    assert missing == [BANDS.index(171), BANDS.index(304)]
    assert render_real.bands_to_fill(cases.Overlay(archive={
        (16, w, DAY0): "/x" for w in BANDS}), set(), 16, DAY0, BANDS) == []


# ------------------------------------------------------------------- resume


def test_done_markers_are_per_day_and_satellite(tmp_path):
    state = render_real.State(str(tmp_path / "out"))
    state.mark_done(DAY0, 16, {"written": 1})
    assert state.is_done(DAY0, 16)
    assert not state.is_done(DAY0, 18) and not state.is_done(DAY0 + DAY, 16)


def test_done_markers_are_per_filler(tmp_path):
    """Two fillers rendered into one root must not skip each other's days."""
    out = str(tmp_path / "out")
    render_real.State(out, "learned").mark_done(DAY0, 16, {})
    assert render_real.State(out, "learned").is_done(DAY0, 16)
    assert not render_real.State(out, "optical_flow").is_done(DAY0, 16)


# -------------------------------------------------------- optical-flow filler


@pytest.fixture
def flow_calls(monkeypatch):
    """Stub the flow itself and the reads: record each fill's anchors."""
    calls = []

    def flow(context):
        calls.append((context.dt_before, context.dt_after))
        if context.before is None and context.after is None:
            return None
        return np.full((4, 4), 5.0, dtype=np.float32)

    monkeypatch.setattr(render_real.bench.fillers, "optical_flow", flow)
    monkeypatch.setattr(render_real.bench.fitsio, "read_image",
                        lambda path: (np.ones((4, 4), dtype=np.float32), None))
    return calls


def test_flow_bridges_from_the_nearest_clean_frames_skipping_bad_ones(flow_calls):
    step = paths.CADENCE
    archive = {(16, 171, DAY0 + i * step): "/x" for i in range(10) if i != 5}
    bad = {(16, 171, DAY0 + 4 * step), (16, 171, DAY0 + 7 * step)}
    fill = render_real.FlowFiller(cases.Overlay(archive=archive), BANDS, bad)
    filled = fill.fill(16, DAY0 + 5 * step, [BANDS.index(171)])
    assert float(filled[BANDS.index(171)][0, 0]) == 5.0
    assert flow_calls == [(2 * step, step)]     # slot 3 (4 is bad), slot 6


def test_flow_fills_only_the_bands_asked_for(flow_calls):
    archive = {(16, w, DAY0 + i * paths.CADENCE): "/x" for w in BANDS for i in (0, 2)}
    fill = render_real.FlowFiller(cases.Overlay(archive=archive), BANDS, set())
    filled = fill.fill(16, DAY0 + paths.CADENCE, [0, 3])
    assert set(filled) == {0, 3} and len(flow_calls) == 2


def test_a_band_with_no_clean_frame_anywhere_is_unfillable(flow_calls):
    archive = {(16, w, DAY0): "/x" for w in BANDS if w != 304}
    fill = render_real.FlowFiller(cases.Overlay(archive=archive), BANDS, set())
    assert fill.fill(16, DAY0 + paths.CADENCE, [BANDS.index(171)]) is not None
    assert fill.fill(16, DAY0 + paths.CADENCE,
                     [BANDS.index(171), BANDS.index(304)]) is None


def test_the_flow_anchor_cache_is_bounded(flow_calls):
    step = paths.CADENCE
    archive = {(16, 171, DAY0 + i * step): "/x" for i in range(0, 400, 2)}
    fill = render_real.FlowFiller(cases.Overlay(archive=archive), BANDS, set())
    for i in range(1, 399, 2):
        fill.fill(16, DAY0 + i * step, [BANDS.index(171)])
        assert len(fill.cache) <= render_real.FLOW_CACHE + 2


# ------------------------------------------------------------- the day loop


class _Filler:
    """Stands in for the model: records what it was asked to fill."""
    asked = []
    result = "array"
    #: Band indices the stand-in model cannot cover (returned all-NaN).
    uncovered = ()

    def __init__(self, overlay, satellites, wavelengths, bad):
        pass

    def fill(self, satellite, when):
        _Filler.asked.append((satellite, when))
        if _Filler.result is None:
            return None
        out = np.full((6, STUB, STUB), 7.0, dtype=np.float32)
        out[list(_Filler.uncovered)] = np.nan
        return out


@pytest.fixture
def day_loop(tmp_path, monkeypatch):
    """A render_day harness with the model, compositor and archive reads stubbed."""
    _Filler.asked, _Filler.result, _Filler.uncovered = [], "array", ()
    flushed = []
    monkeypatch.setattr(render_real.bench, "_LearnedFiller", _Filler)
    monkeypatch.setattr(render_real.bench, "_flush_renders",
                        lambda jobs, workers: flushed.extend(list(jobs)))
    monkeypatch.setattr(render_real.fitsio, "read_image",
                        lambda path: (np.ones((STUB, STUB), dtype=np.float32), None))

    def run(rows, bad=frozenset(), satellites=(16, 18), max_dual_gap=360,
            filler="learned"):
        rows = [r[:5] + ("bad" if r[:3] in bad else "good",) for r in rows]
        state = render_real.State(str(tmp_path / "out"), filler)
        summary = render_real.render_day(
            DAY0, 16, satellites, BANDS, rows, str(tmp_path), state,
            str(tmp_path / "out"), workers=1, batch=1000, max_dual_gap=max_dual_gap,
            filler=filler)
        return summary, flushed

    return run


def full_rows(satellites=(16, 18), skip=(), lacking=None):
    """Index rows for every slot of the window, minus (satellite, time) in skip
    and the bands `lacking` maps a satellite to (absent all window)."""
    t0, t1 = render_real.window_for(DAY0)
    lacking = lacking or {}
    return [(s, w, t, f"{s}_{w}_{t}.fits", 1_500_000, "good")
            for s in satellites for t in cases.timeline(t0, t1) for w in BANDS
            if (s, t) not in skip and w not in lacking.get(s, ())]


# ------------------------------------------------- bands split across satellites


def test_a_day_with_each_band_on_one_satellite_or_the_other_is_rendered(day_loop):
    """2024-03-03: GOES-16 had no 94/131, GOES-18 no 304, all day.  Neither held a
    complete frame, yet every band was observed -- it is not a dual outage."""
    summary, _ = day_loop(full_rows(lacking={16: (94, 131), 18: (304,)}),
                          max_dual_gap=30)
    assert summary["policy_gap"] == 0
    assert summary["written"] == 360 and summary["filled"] == 360
    assert summary["bad_bands"] == 2 * 360


def test_a_band_neither_satellite_holds_stays_black(day_loop):
    summary, _ = day_loop(full_rows(lacking={16: (94,), 18: (94,)}), max_dual_gap=30)
    assert summary["policy_gap"] == 360 and summary["written"] == 0
    assert _Filler.asked == []


def test_a_needed_band_the_model_cannot_cover_is_unfillable(day_loop):
    """NaN means no frame in the stack held the band: never composite it."""
    target = DAY0 + 10 * paths.CADENCE
    _Filler.uncovered = (BANDS.index(171),)
    summary, flushed = day_loop(full_rows(), bad={(16, 171, target)})
    assert summary["unfillable"] == 1
    assert target not in {when for _, when, _ in flushed}


def test_an_uncovered_band_the_slot_does_not_need_is_harmless(day_loop):
    target = DAY0 + 10 * paths.CADENCE
    _Filler.uncovered = (BANDS.index(304),)
    summary, flushed = day_loop(full_rows(), bad={(16, 171, target)})
    assert summary["unfillable"] == 0 and summary["filled"] == 1
    arrays = next(a for a, when, _ in flushed if when == target)
    assert all(np.isfinite(band).all() for band in arrays)


def test_unfillable_by_band_unions_the_bands_long_gaps():
    step = paths.CADENCE
    timeline = list(range(DAY0, DAY0 + 100 * step, step))
    archive = {(s, w, t): "/x" for s in (16, 18) for w in BANDS for t in timeline}
    # 94 unobserved by both for 40 slots; 304 by both for 10; 131 split.
    for t in timeline[10:50]:
        del archive[(16, 94, t)], archive[(18, 94, t)]
    for t in timeline[60:70]:
        del archive[(16, 304, t)], archive[(18, 304, t)]
    for t in timeline:
        del archive[(16, 131, t)]
    bad = {(18, 131, t) for t in timeline[80:90]}           # bad counts as unobserved
    overlay = cases.Overlay(archive=archive)
    from bench import unfillable_by_band
    assert unfillable_by_band(overlay, (16, 18), BANDS, set(), timeline, 30) == \
        set(timeline[10:50])
    assert unfillable_by_band(overlay, (16, 18), BANDS, bad, timeline, 5) == \
        set(timeline[10:50]) | set(timeline[60:70]) | set(timeline[80:90])


def test_a_clean_day_is_composited_whole_without_the_model(day_loop):
    summary, flushed = day_loop(full_rows())
    assert summary["written"] == DAY // paths.CADENCE == 360
    assert summary["filled"] == 0 and _Filler.asked == []
    assert len({when for _, when, _ in flushed}) == 360


def test_only_the_bad_bands_come_from_the_model(day_loop):
    target = DAY0 + 100 * paths.CADENCE
    summary, flushed = day_loop(full_rows(), bad={(16, 171, target)})
    assert _Filler.asked == [(16, target)]
    arrays = next(a for a, when, _ in flushed if when == target)
    assert float(arrays[BANDS.index(171)][0, 0]) == 7.0     # from the model
    assert float(arrays[BANDS.index(94)][0, 0]) == 1.0      # from the archive
    assert summary["filled"] == 1 and summary["bad_bands"] == 1


def test_a_slot_absent_from_the_archive_is_filled_too(day_loop):
    """Every cadence slot is composited, not just those the archive holds."""
    target = DAY0 + 50 * paths.CADENCE
    summary, _ = day_loop(full_rows(skip={(16, target)}))
    assert (16, target) in _Filler.asked
    assert summary["written"] == 360


def test_a_dual_gap_beyond_policy_stays_black(day_loop):
    """Both satellites dark for longer than the policy: nothing is invented."""
    gap = {(s, DAY0 + i * paths.CADENCE) for s in (16, 18) for i in range(40)}
    summary, flushed = day_loop(full_rows(skip=gap), max_dual_gap=30)
    assert summary["policy_gap"] == 40
    assert summary["written"] == 320
    assert not any((16, DAY0 + i * paths.CADENCE) in _Filler.asked for i in range(40))


def test_a_short_dual_gap_within_policy_is_filled(day_loop):
    gap = {(s, DAY0 + i * paths.CADENCE) for s in (16, 18) for i in range(10)}
    summary, _ = day_loop(full_rows(skip=gap), max_dual_gap=30)
    assert summary["policy_gap"] == 0 and summary["filled"] == 10


def test_nothing_usable_in_the_stack_skips_rather_than_inventing(day_loop):
    _Filler.result = None
    target = DAY0 + 7 * paths.CADENCE
    summary, flushed = day_loop(full_rows(), bad={(16, 94, target)})
    assert summary["unfillable"] == 1
    assert target not in {when for _, when, _ in flushed}


def test_existing_composites_are_never_rewritten(day_loop, tmp_path):
    out_dir = tmp_path / "out" / "learned" / "goes16"
    out_dir.mkdir(parents=True)
    (out_dir / f"Composite-{DAY0}.jpg").write_bytes(b"done")
    summary, flushed = day_loop(full_rows())
    assert summary["existing"] == 1 and summary["written"] == 359
    assert (out_dir / f"Composite-{DAY0}.jpg").read_bytes() == b"done"


def test_the_optical_flow_filler_repairs_only_the_bad_bands(day_loop, monkeypatch,
                                                           tmp_path):
    target = DAY0 + 100 * paths.CADENCE
    asked = []

    def replacement(variant, overlay, slot, bad, cache, wavelengths):
        asked.append((variant, slot, slot in bad))
        return np.full((STUB, STUB), 3.0, dtype=np.float32), True

    monkeypatch.setattr(render_real.bench, "_replacement_band", replacement)
    summary, flushed = day_loop(full_rows(), bad={(16, 171, target)},
                                filler="optical_flow")
    assert asked == [("new", (16, 171, target), True)] and _Filler.asked == []
    arrays = next(a for a, when, _ in flushed if when == target)
    assert float(arrays[BANDS.index(171)][0, 0]) == 3.0
    assert float(arrays[BANDS.index(94)][0, 0]) == 1.0
    assert summary["filled"] == 1 and summary["written"] == 360
    assert all("/optical_flow/goes16/" in path for _, _, path in flushed)


def test_a_day_with_unjudged_frames_is_refused(tmp_path):
    """Rendering on a weaker test is how the bars reached the first 2024 video."""
    rows = full_rows()
    rows[0] = rows[0][:5] + (None,)
    out = str(tmp_path / "out")
    with pytest.raises(SystemExit, match="no stored verdict"):
        render_real.render_day(DAY0, 16, (16, 18), BANDS, rows, str(tmp_path),
                               render_real.State(out), out, 1, 1000, 360)


def test_no_judgment_run_under_the_current_rules_is_refused(monkeypatch, tmp_path):
    from suvi import db

    db.connect(str(tmp_path / "index.sqlite")).close()
    monkeypatch.setattr(render_real.fillers, "load_learned", lambda: ("net",))
    with pytest.raises(SystemExit, match="No production judgment"):
        render_real.main(["--start", "2024-06-01", "--end", "2024-06-02",
                          "--out", str(tmp_path / "out"), "--db", str(tmp_path / "index.sqlite")])


def test_optical_flow_needs_no_checkpoint(monkeypatch, tmp_path):
    monkeypatch.setattr(render_real.fillers, "load_learned", lambda: None)
    monkeypatch.setattr(render_real.db, "connect",
                        lambda *a, **k: (_ for _ in ()).throw(SystemExit("reached the index")))
    with pytest.raises(SystemExit, match="reached the index"):
        render_real.main(["--start", "2024-06-01", "--end", "2024-06-02",
                          "--out", str(tmp_path), "--filler", "optical_flow"])
    with pytest.raises(SystemExit, match="needs a checkpoint"):
        render_real.main(["--start", "2024-06-01", "--end", "2024-06-02",
                          "--out", str(tmp_path)])


def test_the_renderer_opens_the_index_read_only(monkeypatch, tmp_path):
    """The downloader writes the index concurrently; a second writer is what
    corrupted it before."""
    seen = {}

    def connect(path=None, readonly=False, timeout=60.0):
        seen["readonly"] = readonly
        raise SystemExit("stop here")

    monkeypatch.setattr(render_real.db, "connect", connect)
    monkeypatch.setattr(render_real.fillers, "load_learned", lambda: ("net",))
    with pytest.raises(SystemExit, match="stop here"):
        render_real.main(["--start", "2024-06-01", "--end", "2024-06-02",
                          "--out", str(tmp_path)])
    assert seen["readonly"] is True