noaa-goes-visualization/tests/test_comparison_video.py

"""Tests for assembling the three-pane comparison video.

The property that matters is *alignment*: a variant that produced no composite for a
timestamp -- which is exactly what the current pipeline does past its gap limit --
must still occupy that slot, or its pane slides out of step with the others and the
comparison is meaningless.
"""

import os

import numpy as np
from PIL import Image

import make_comparison_video as mcv
from suvi import paths

T0 = 1715299200


def write_stream(root, variant, satellite, timestamps, size=(64, 36)):
    directory = os.path.join(root, variant, f"goes{satellite}")
    os.makedirs(directory, exist_ok=True)
    for when in timestamps:
        img = Image.fromarray(np.full((size[1], size[0], 3), 40, dtype="uint8"))
        img.save(os.path.join(directory, f"Composite-{when}.jpg"), quality=90)
    return directory


def test_stream_frames_reads_timestamps(tmp_path):
    times = [T0, T0 + paths.CADENCE, T0 + 2 * paths.CADENCE]
    directory = write_stream(str(tmp_path), "pristine", 16, times)
    frames = mcv.stream_frames(directory)
    assert sorted(frames) == times
    assert all(os.path.exists(p) for p in frames.values())


def test_stream_frames_ignores_other_files(tmp_path):
    directory = write_stream(str(tmp_path), "new", 16, [T0])
    open(os.path.join(directory, "notes.txt"), "w").write("x")
    open(os.path.join(directory, "Composite-nonsense.jpg"), "w").write("x")
    assert list(mcv.stream_frames(directory)) == [T0]


def test_stream_frames_on_a_missing_directory(tmp_path):
    assert mcv.stream_frames(str(tmp_path / "absent")) == {}


def test_every_variant_is_labelled():
    assert set(mcv.LABELS) == set(mcv.VARIANTS)
    assert all(mcv.LABELS[v] for v in mcv.VARIANTS)


def test_encode_emits_one_frame_per_slot_including_gaps(tmp_path, monkeypatch):
    """A pane with holes must still be full length, or the panes drift apart.

    'today' legitimately has none for long gaps; those slots have to become black
    frames occupying the timeline, not vanish from it.
    """
    times = [T0, T0 + 2 * paths.CADENCE]  # slot 1 deliberately missing
    directory = write_stream(str(tmp_path), "today", 16, times)
    frames = mcv.stream_frames(directory)
    timeline = [T0 + i * paths.CADENCE for i in range(3)]

    written = []

    class FakeProcess:
        def __init__(self):
            self.stdin = self

        def write(self, data):
            written.append(("real", len(data)))

        def save(self, *a, **k):
            pass

        def close(self):
            pass

        def wait(self):
            return 0

    fake = FakeProcess()
    monkeypatch.setattr(mcv.subprocess, "Popen", lambda *a, **k: fake)

    class CountingImage:
        size = (64, 36)

        @staticmethod
        def fromarray(arr):
            class Blank:
                def save(self, handle, *a, **k):
                    written.append(("black", 0))
            return Blank()

        @staticmethod
        def open(path):
            return CountingImage

    monkeypatch.setattr(mcv, "Image", CountingImage)

    present = mcv.encode_stream("ffmpeg", frames, timeline, "out.mp4", 60, 18)
    assert present == 2
    assert len(written) == 3, "a missing slot did not occupy the timeline"
    assert [kind for kind, _ in written] == ["real", "black", "real"]


# ------------------------------------------------------ pristine reads the truth


def test_pristine_resolves_deleted_slots_from_the_archive(tmp_path):
    """The ground-truth pane must show the window as it really is.

    A slot the case deleted has no overlay path.  Resolving pristine through the
    overlay silently drops those timestamps -- 221 of them in one run -- leaving the
    reference pane shorter than the panes it exists to be compared against.
    """
    import bench
    from suvi import cases

    slot = (16, 171, T0)
    overlay = cases.Overlay(
        archive={slot: ("/archive/real.fits", None)},
        overrides={},
        deleted=frozenset({slot}),
    )
    assert overlay.path(slot) is None, "precondition: the case deleted this slot"
    assert overlay.truth_path(slot) == "/archive/real.fits"

    reads = []

    def fake_read(path):
        reads.append(path)
        return "pixels", None

    original = bench.fitsio.read_image
    bench.fitsio.read_image = fake_read
    try:
        cache = {}
        assert bench._resolve_band(overlay, slot, cache, use_truth=True) == "pixels"
        assert reads == ["/archive/real.fits"]
        # Without use_truth the slot is genuinely gone, which is right for the
        # variants that are meant to show the damage.
        assert bench._resolve_band(overlay, slot, cache, use_truth=False) is None
    finally:
        bench.fitsio.read_image = original


def test_resolve_band_caches_truth_and_overlay_separately(tmp_path):
    """One cache serving both would return the wrong frame for one of them."""
    import bench
    from suvi import cases

    slot = (16, 171, T0)
    overlay = cases.Overlay(
        archive={slot: ("/archive/real.fits", None)},
        overrides={slot: "/overlay/damaged.fits"},
    )
    original = bench.fitsio.read_image
    bench.fitsio.read_image = lambda path: (path, None)
    try:
        cache = {}
        assert bench._resolve_band(overlay, slot, cache, use_truth=True) == "/archive/real.fits"
        assert bench._resolve_band(overlay, slot, cache, use_truth=False) == "/overlay/damaged.fits"
    finally:
        bench.fitsio.read_image = original