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