noaa-goes-visualization/tests/test_cases.py

import os

import pytest

from conftest import solar_disc, write_fits
from suvi import cases, corruptions, paths

T0 = 1715299200  # 2024-05-10 00:00 UTC
SATS = (16, 18)
BANDS = (171, 304)


def populate(root, slots, label="f"):
    """Write minimal frames for the given slots into a fake archive."""
    written = {}
    for satellite, wavelength, when in slots:
        name = paths.FrameName(satellite, wavelength, when, when + paths.CADENCE, "1-0-2")
        path = os.path.join(str(root), *name.relpath(label).split("/"))
        write_fits(path, solar_disc(size=64, radius=20, peak=3.0))
        written[name.slot] = path
    return written


def grid(count, satellites=SATS, wavelengths=BANDS, start=T0):
    return [
        (satellite, wavelength, start + index * paths.CADENCE)
        for index in range(count)
        for satellite in satellites
        for wavelength in wavelengths
    ]


# ----------------------------------------------------------------------- scanning


def test_scan_window_finds_frames_and_labels(archive):
    populate(archive, grid(3))
    found = cases.scan_window(str(archive), SATS, BANDS, T0, T0 + 3 * paths.CADENCE)
    assert len(found) == 12
    path, label = found[(16, 171, T0)]
    assert label == "f" and os.path.exists(path)


def test_scan_window_respects_the_time_bounds(archive):
    populate(archive, grid(5))
    found = cases.scan_window(str(archive), SATS, BANDS, T0, T0 + 2 * paths.CADENCE)
    assert {slot[2] for slot in found} == {T0, T0 + paths.CADENCE}


def test_scan_window_ignores_unrelated_files(archive):
    populate(archive, grid(1))
    stray = os.path.join(str(archive), "goes16/l2/data/suvi-l2-ci171/2024/05/10/notes.txt")
    open(stray, "w").write("not a frame")
    found = cases.scan_window(str(archive), SATS, BANDS, T0, T0 + paths.CADENCE)
    assert len(found) == 4


def test_scan_window_handles_a_missing_archive(tmp_path):
    assert cases.scan_window(str(tmp_path), SATS, BANDS, T0, T0 + 240) == {}


def test_scan_window_spans_a_day_boundary(archive):
    midnight = 1715299200 - 2 * paths.CADENCE  # last slots of the previous day
    populate(archive, [(16, 171, midnight + i * paths.CADENCE) for i in range(4)])
    found = cases.scan_window(str(archive), (16,), (171,), midnight, midnight + 4 * 240)
    assert len(found) == 4


# --------------------------------------------------------------------------- runs


def test_timeline_covers_the_window():
    assert cases.timeline(T0, T0 + 3 * paths.CADENCE) == [
        T0, T0 + paths.CADENCE, T0 + 2 * paths.CADENCE
    ]


def test_find_runs_reports_the_longest_first(archive):
    slots = grid(10)
    # Punch a hole at index 4 for one band on one satellite.
    slots = [s for s in slots if s != (16, 171, T0 + 4 * paths.CADENCE)]
    populate(archive, slots)
    found = cases.scan_window(str(archive), SATS, BANDS, T0, T0 + 10 * paths.CADENCE)
    runs = cases.find_runs(found, SATS, BANDS, T0, T0 + 10 * paths.CADENCE, minimum=2)
    assert runs[0][0] == 5  # indices 5..9
    assert runs[1][0] == 4  # indices 0..3


def test_find_runs_ignores_labels_by_default(archive):
    """Labels are stale, so completeness is what select-window reports."""
    populate(archive, grid(6), label="e")  # everything marked bad by the old filter
    found = cases.scan_window(str(archive), SATS, BANDS, T0, T0 + 6 * paths.CADENCE)
    assert cases.find_runs(found, SATS, BANDS, T0, T0 + 6 * paths.CADENCE, minimum=2)
    assert not cases.find_runs(
        found, SATS, BANDS, T0, T0 + 6 * paths.CADENCE, minimum=2, require_good_label=True
    )


def test_slot_is_good_treats_unlabelled_as_unknown(archive):
    populate(archive, grid(1), label=None)
    found = cases.scan_window(str(archive), SATS, BANDS, T0, T0 + paths.CADENCE)
    assert cases.slot_is_good(found, 16, 171, T0) is None
    assert cases.slot_is_good(found, 16, 171, T0 + 9999) is None


# --------------------------------------------------------------------------- plan


def test_plan_validates_its_inputs():
    cases.InjectionPlan().validate()
    with pytest.raises(ValueError, match="fraction"):
        cases.InjectionPlan(fraction=1.5).validate()
    with pytest.raises(ValueError, match="gap_lengths"):
        cases.InjectionPlan(gap_lengths=(0,)).validate()
    with pytest.raises(ValueError, match="satellite_scope"):
        cases.InjectionPlan(satellite_scope="g99").validate()
    with pytest.raises(ValueError, match="wavelength_scope"):
        cases.InjectionPlan(wavelength_scope="some").validate()
    with pytest.raises(ValueError, match="unknown corruption"):
        cases.InjectionPlan(modes=("nonsense",)).validate()
    with pytest.raises(ValueError, match="severity"):
        cases.InjectionPlan(severity=(0.9, 0.1)).validate()


def test_plan_round_trips_through_a_dict():
    plan = cases.InjectionPlan(fraction=0.2, satellite_scope="both", modes=("delete",))
    assert cases.InjectionPlan.from_dict(plan.as_dict()) == plan


def test_injections_are_deterministic():
    slots = set(grid(200))
    first = cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=7)
    second = cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=7)
    assert first == second


def test_a_different_seed_gives_a_different_plan():
    slots = set(grid(200))
    first = cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=1)
    second = cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=2)
    assert first != second


def test_injections_leave_the_window_edges_intact():
    """Every damaged slot needs good frames either side to be reconstructed from."""
    slots = set(grid(200))
    times = sorted({slot[2] for slot in slots})
    injections = cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=3)
    damaged = {i.slot[2] for i in injections}
    assert damaged.isdisjoint(times[: cases.EDGE_MARGIN])
    assert damaged.isdisjoint(times[-cases.EDGE_MARGIN :])


def test_gaps_stay_separated():
    slots = set(grid(300))
    times = sorted({slot[2] for slot in slots})
    index_of = {when: i for i, when in enumerate(times)}
    injections = cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=5)
    damaged = sorted(index_of[i.slot[2]] for i in {(j.slot[2],): j for j in injections}.values())
    runs = []
    for position in damaged:
        if runs and position == runs[-1][-1] + 1:
            runs[-1].append(position)
        else:
            runs.append([position])
    for earlier, later in zip(runs, runs[1:]):
        assert later[0] - earlier[-1] > cases.MIN_SEPARATION


def test_gap_metadata_is_recorded():
    slots = set(grid(200))
    injections = cases.plan_injections(
        cases.InjectionPlan(gap_lengths=(5,)), slots, SATS, BANDS, seed=9
    )
    assert injections
    for injection in injections:
        assert injection.gap_length == 5
        assert 0 <= injection.gap_index < 5
        assert injection.mode == "delete" or injection.mode in corruptions.CATALOG


def test_satellite_scope_both_damages_each_satellite_together():
    slots = set(grid(200))
    injections = cases.plan_injections(
        cases.InjectionPlan(satellite_scope="both"), slots, SATS, BANDS, seed=11
    )
    by_time = {}
    for injection in injections:
        by_time.setdefault(injection.slot[2], set()).add(injection.slot[0])
    assert all(sats == set(SATS) for sats in by_time.values())


def test_satellite_scope_can_target_one_spacecraft():
    slots = set(grid(200))
    injections = cases.plan_injections(
        cases.InjectionPlan(satellite_scope="g18"), slots, SATS, BANDS, seed=11
    )
    assert {i.slot[0] for i in injections} == {18}


def test_wavelength_scope_one_damages_a_single_band():
    slots = set(grid(200))
    injections = cases.plan_injections(
        cases.InjectionPlan(wavelength_scope="one"), slots, SATS, BANDS, seed=13
    )
    by_time = {}
    for injection in injections:
        by_time.setdefault(injection.slot[2], set()).add(injection.slot[1])
    assert all(len(bands) == 1 for bands in by_time.values())


def test_zero_fraction_plans_nothing():
    slots = set(grid(200))
    assert cases.plan_injections(
        cases.InjectionPlan(fraction=0.0), slots, SATS, BANDS, seed=1
    ) == []


def test_a_window_too_short_to_damage():
    slots = set(grid(4))
    assert cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=1) == []


# ------------------------------------------------------------------------ overlay


def build_overlay():
    archive = {(16, 171, T0 + i * 240): (f"/archive/{i}.fits", "f") for i in range(5)}
    return cases.Overlay(
        archive=archive,
        overrides={(16, 171, T0 + 240): "/overlay/1.fits"},
        deleted=frozenset({(16, 171, T0 + 480)}),
    )


def test_overlay_passes_untouched_slots_through():
    assert build_overlay().path((16, 171, T0)) == "/archive/0.fits"


def test_overlay_redirects_corrupted_slots():
    assert build_overlay().path((16, 171, T0 + 240)) == "/overlay/1.fits"


def test_overlay_hides_deleted_slots():
    assert build_overlay().path((16, 171, T0 + 480)) is None


def test_overlay_always_exposes_the_original_for_scoring():
    overlay = build_overlay()
    assert overlay.truth_path((16, 171, T0 + 240)) == "/archive/1.fits"
    assert overlay.truth_path((16, 171, T0 + 480)) == "/archive/2.fits"


def test_overlay_reports_an_unknown_slot_as_absent():
    assert build_overlay().path((99, 999, 0)) is None


def test_overlay_series_is_ordered_and_filtered():
    overlay = cases.Overlay(
        archive={
            (16, 171, T0 + 480): ("/c", None),
            (16, 171, T0): ("/a", None),
            (18, 171, T0 + 240): ("/b", None),
        }
    )
    assert overlay.series(16, 171) == [(16, 171, T0), (16, 171, T0 + 480)]
    assert overlay.series(18, 171) == [(18, 171, T0 + 240)]


# ------------------------------------------------------------------ long gaps


def test_separation_scales_with_gap_length():
    """Four clean slots is enough beside a 2-frame gap and useless beside a 300."""
    assert cases.gap_separation(1) == cases.MIN_SEPARATION
    assert cases.gap_separation(6) == cases.MIN_SEPARATION
    assert cases.gap_separation(300) == 75
    assert cases.gap_separation(100) == 25


def test_long_gaps_keep_clean_brackets():
    slots = set(grid(1200))
    times = sorted({s[2] for s in slots})
    index_of = {t: i for i, t in enumerate(times)}
    injections = cases.plan_injections(
        cases.InjectionPlan(fraction=0.3, gap_lengths=(3, 30, 100)),
        slots, SATS, BANDS, seed=5,
    )
    damaged = sorted({index_of[i.slot[2]] for i in injections})

    runs = []
    for position in damaged:
        if runs and position == runs[-1][-1] + 1:
            runs[-1].append(position)
        else:
            runs.append([position])
    for earlier, later in zip(runs, runs[1:]):
        needed = cases.gap_separation(len(earlier))
        assert later[0] - earlier[-1] > needed, (
            f"a {len(earlier)}-slot gap had only {later[0] - earlier[-1]} clean "
            f"slots before the next, needs more than {needed}"
        )


def test_a_window_too_short_for_the_gaps_is_refused():
    """Silently planning nothing would report results for lengths never tested."""
    slots = set(grid(50))
    with pytest.raises(ValueError, match="usable slots"):
        cases.plan_injections(
            cases.InjectionPlan(gap_lengths=(300,)), slots, SATS, BANDS, seed=1
        )


def test_long_gaps_actually_get_placed():
    """Longest-first placement: a 300 laid down last would rarely find room."""
    slots = set(grid(2600))
    injections = cases.plan_injections(
        cases.InjectionPlan(fraction=0.25, gap_lengths=(1, 2, 3, 10, 30, 100, 300)),
        slots, SATS, BANDS, seed=7,
    )
    placed = {i.gap_length for i in injections}
    assert 300 in placed, "the longest gap never got placed"
    assert placed & {1, 2, 3}, "short gaps were crowded out"


def test_gap_lengths_span_short_and_long():
    slots = set(grid(2600))
    injections = cases.plan_injections(
        cases.InjectionPlan(fraction=0.25, gap_lengths=(1, 2, 3, 10, 30, 100, 300)),
        slots, SATS, BANDS, seed=3,
    )
    lengths = {i.gap_length for i in injections}
    assert min(lengths) <= 3 and max(lengths) >= 100


def test_gaps_per_length_places_the_requested_count():
    """A fraction budget is eaten by the longest gaps, leaving one of each.

    An earlier case ended up 84% a single 300-slot gap of a single mode, which
    measures that mode at that length and nothing else.
    """
    slots = set(grid(2600))
    injections = cases.plan_injections(
        cases.InjectionPlan(gap_lengths=(1, 3, 10), gaps_per_length=4),
        slots, SATS, BANDS, seed=4,
    )
    counts = {}
    for injection in injections:
        key = (injection.gap_length, injection.slot[2] - injection.gap_index * paths.CADENCE)
        counts.setdefault(injection.gap_length, set()).add(key[1])
    assert counts[1] == counts[1] and len(counts[1]) == 4
    assert len(counts[3]) == 4
    assert len(counts[10]) == 4


def test_gaps_per_length_overrides_the_fraction():
    slots = set(grid(2600))
    injections = cases.plan_injections(
        cases.InjectionPlan(fraction=0.001, gap_lengths=(10,), gaps_per_length=3),
        slots, SATS, BANDS, seed=6,
    )
    starts = {i.slot[2] - i.gap_index * paths.CADENCE for i in injections}
    assert len(starts) == 3, "the tiny fraction should have been ignored"


def test_modes_are_cycled_not_resampled():
    """With few gaps, independent draws repeat and leave the catalogue untested."""
    slots = set(grid(2600))
    injections = cases.plan_injections(
        cases.InjectionPlan(gap_lengths=(2,), gaps_per_length=8), slots, SATS, BANDS, seed=8
    )
    starts = {}
    for injection in injections:
        starts.setdefault(injection.slot[2] - injection.gap_index * paths.CADENCE,
                          injection.mode)
    assert len(set(starts.values())) == len(starts), "a mode was reused across gaps"


def test_plan_with_gaps_per_length_round_trips():
    plan = cases.InjectionPlan(gaps_per_length=5)
    assert cases.InjectionPlan.from_dict(plan.as_dict()).gaps_per_length == 5