scanengine-3/tests/test_scan_engine.py

"""Headless ScanEngine tests driven entirely by fake hardware.

These cover what can't be checked without the rig: the command sequence,
the written file layout, and abort/pause behaviour.
"""
import threading
import time

import pytest

from core.rotation import RotationAxis, RotationSettings
from core.scan_engine import (
    AXIS_X, AXIS_Y, ScanAborted, ScanCallbacks, ScanEngine, ResumeState,
    ResumeTarget,
)
from core.scan_geometry import ScanGeometryError, build_plan
from core.sras_format import BG_LEN_SIZE, SCAN_CHANNELS, SrasFile
from fakes import FakeScope, FakeStage, FakeT3R, Trace, background_record

SPF = 8


def make_plan(num_angles=1, y_delta=0.005):
    # Small ROI well inside the stage limits: few rows, few frames per angle.
    return build_plan(40.0, 30.0, 0.02, y_delta, num_angles, 0.01,
                      laser_freq_hz=20000.0, velocity_mm_s=100.0)


def build(tmp_path, num_angles=1, callbacks=None, resume=None, plan=None,
          burst_mode=False, max_frames=4096, out_name="out.sras",
          strict_rows=False, **kw):
    trace = Trace()
    scope = FakeScope(trace, samples_per_frame=SPF, max_frames=max_frames)
    stage = FakeStage(trace, scope=scope)
    t3r = FakeT3R(trace, **kw)
    rotator = RotationAxis(t3r, RotationSettings())
    plan = plan if plan is not None else make_plan(num_angles)
    engine = ScanEngine(stage, scope, rotator, plan, tmp_path / out_name,
                        resume=resume,
                        callbacks=callbacks or ScanCallbacks(),
                        burst_mode=burst_mode, strict_rows=strict_rows)
    return engine, trace, plan


def test_single_angle_scan_writes_readable_file(tmp_path):
    engine, trace, plan = build(tmp_path)
    result = engine.run()

    assert not result.aborted
    assert result.rows_written == plan.per_angle[0].n_rows
    assert result.angles_acquired == [0]

    sras = SrasFile(result.path)
    assert sras.header.n_angles == 1
    assert sras.header.samples_per_frame == SPF
    assert sras.header.n_channels == len(SCAN_CHANNELS)
    # File is complete: every declared row present on disk
    assert [s.status for s in sras.angle_status()] == ["OK"]
    assert len(sras.preambles) == 3
    assert sras.backgrounds == [background_record(0, SPF)]


def test_each_angle_captures_and_stores_its_own_background(tmp_path):
    """One background per angle, taken after the rotation, kept ahead of it."""
    prompts = []
    engine, trace, plan = build(
        tmp_path, num_angles=3,
        callbacks=ScanCallbacks(prompt=lambda title, msg: prompts.append(title)))
    engine.run()

    sras = SrasFile(tmp_path / "out.sras")
    assert trace.count("transfer_curve") == 3
    assert sras.backgrounds == [background_record(i, SPF) for i in range(3)]
    for st in sras.angle_status():
        assert st.bg_bytes == BG_LEN_SIZE + SPF
        assert st.status == "OK"

    # Each capture follows the rotation to the angle it belongs to (the last
    # rotation is the return to home after the final angle).
    assert [c[0] for c in trace.calls
            if c[0] in ("t3r_rotate", "transfer_curve")] == [
        "transfer_curve", "t3r_rotate", "transfer_curve",
        "t3r_rotate", "transfer_curve", "t3r_rotate"]

    # Two prompts per angle: Genesis off for the capture, back on to scan.
    assert prompts == [title for i in range(3) for title in
                       (f"Background Capture — Angle {i + 1}/3",
                        f"Begin Angle {i + 1}/3")]

    # The scope goes back to the scan trigger after every capture, not just
    # once at the start — the capture needs the single-record edge trigger.
    cmds = [c[1] for c in trace.of("write")]
    assert cmds.count("TRIGger:A:TYPe EDGE") == 4      # prepare + one per angle
    assert cmds.count("TRIGger:A:TYPe LOGIc") == 4


def test_command_sequence_order(tmp_path):
    engine, trace, plan = build(tmp_path)
    engine.run()
    names = trace.names()

    def first(name):
        return names.index(name)

    # Stage prepared, then scope configured, then rows executed
    assert first("set_trigger_trigout_maxv") < first("set_sample_rate")
    assert first("set_sample_rate") < first("transfer_fastframe")
    # Velocity set for both axes before any scan move
    assert trace.count("set_velocity_params") == 2
    # Per row: Y positioned, then X pre-ramp, then X run
    moves = trace.of("move_axis_absolute")
    assert moves[0][1] == AXIS_Y
    assert moves[1][1] == AXIS_X and moves[2][1] == AXIS_X
    assert moves[1][2] < moves[2][2]          # pre-ramp start < run-off end
    # Data channels transferred (CH3 is synthesized, not read)
    assert [c[1] for c in trace.of("transfer_fastframe")] == [1, 4]


def test_multi_angle_rotates_and_returns_home(tmp_path):
    engine, trace, plan = build(tmp_path, num_angles=3)
    engine.run()

    rotations = [c[1] for c in trace.of("t3r_rotate")]
    # Three angles at 0/-90/-180 → two moves out, then one back to 0
    assert rotations == [-90.0, -90.0, 180.0]
    # Every move waits for completion instead of sleeping a guess
    assert trace.count("t3r_wait_motion_done") == len(rotations)
    # GR configured once, before any rotation
    assert trace.names().index("t3r_set_microstep") < trace.names().index("t3r_rotate")

    sras = SrasFile(tmp_path / "out.sras")
    assert [s.status for s in sras.angle_status()] == ["OK"] * 3


def test_fastframe_count_rearmed_per_angle(tmp_path):
    engine, trace, plan = build(tmp_path, num_angles=3)
    engine.run()
    counts = [c[1] for c in trace.of("set_fastframe_count")]
    assert counts == [pa.n_frames for pa in plan.per_angle]


def test_abort_before_start_raises_and_stops_early(tmp_path):
    engine, trace, _ = build(tmp_path)
    engine.abort()
    with pytest.raises(ScanAborted):
        engine.run()
    assert trace.count("transfer_fastframe") == 0


def test_abort_during_prompt_unblocks(tmp_path):
    """A prompt that never returns must not deadlock an aborting scan."""
    released = threading.Event()

    def prompt(title, msg):
        # Simulates the GUI bridge: waits until abort flips the flag.
        while not engine.aborted:
            if released.wait(0.01):
                return

    engine, trace, _ = build(tmp_path, callbacks=ScanCallbacks(prompt=prompt))

    errors = []

    def run():
        try:
            engine.run()
        except ScanAborted:
            errors.append("aborted")

    t = threading.Thread(target=run, daemon=True)
    t.start()
    time.sleep(0.2)   # let it reach the first prompt
    engine.abort()
    t.join(timeout=5)
    assert not t.is_alive(), "engine deadlocked on a prompt during abort"
    assert errors == ["aborted"]


def test_pause_and_resume_at_row_boundary(tmp_path):
    states = []
    engine, trace, plan = build(
        tmp_path, num_angles=1,
        callbacks=ScanCallbacks(on_paused_changed=states.append))
    engine.pause()

    done = threading.Event()

    def run():
        try:
            engine.run()
        except ScanAborted:
            pass          # only reachable via the failure escape hatch below
        finally:
            done.set()

    t = threading.Thread(target=run, daemon=True)
    t.start()
    try:
        # The engine's instrument-settling sleeps run before the first row,
        # so poll for the pause rather than assuming a fixed delay.
        deadline = time.monotonic() + 10.0
        while not states and time.monotonic() < deadline:
            time.sleep(0.05)
        paused = bool(states)
        assert paused and states[0] is True, "engine did not report the pause"
    finally:
        # Always release the scan thread; if the pause never arrived, abort
        # too, so a failed assertion can't leave it parked forever.
        if not states:
            engine.abort()
        engine.resume()
    t.join(timeout=10)
    assert done.is_set()
    assert states[-1] is False


def test_dc_bias_callback_reports_per_frame_means(tmp_path):
    rows = []
    engine, trace, plan = build(
        tmp_path, callbacks=ScanCallbacks(on_dc_bias=lambda r, m: rows.append((r, m))))
    engine.run()

    assert len(rows) == plan.per_angle[0].n_rows
    row_idx, means = rows[0]
    assert row_idx == 1
    assert len(means) == plan.per_angle[0].n_frames
    assert all(isinstance(v, float) for v in means)


def test_offstage_plan_rejected_before_touching_hardware(tmp_path):
    trace = Trace()
    scope = FakeScope(trace, samples_per_frame=SPF)
    stage = FakeStage(trace, scope=scope)
    # X range that runs off the 110 mm stage once ramps are added
    plan = build_plan(80.0, 30.0, 40.0, 5.0, 1, 0.25,
                      laser_freq_hz=20000.0, velocity_mm_s=100.0)
    engine = ScanEngine(stage, scope, None, plan, tmp_path / "bad.sras")
    with pytest.raises(ScanGeometryError):
        engine.run()
    assert trace.calls == [], "hardware touched despite invalid geometry"


def test_multi_angle_without_rotator_raises(tmp_path):
    trace = Trace()
    engine = ScanEngine(FakeStage(trace), FakeScope(trace, samples_per_frame=SPF),
                        None, make_plan(3), tmp_path / "x.sras")
    with pytest.raises(RuntimeError, match="T3R rotation stage"):
        engine.run()


def test_missing_hardware_raises(tmp_path):
    trace = Trace()
    with pytest.raises(RuntimeError, match="BBD202"):
        ScanEngine(None, FakeScope(trace), None, make_plan(),
                   tmp_path / "x.sras").run()
    with pytest.raises(RuntimeError, match="Oscilloscope"):
        ScanEngine(FakeStage(trace), None, None, make_plan(),
                   tmp_path / "x.sras").run()


def test_resume_seeks_to_angle_offset_and_skips_others(tmp_path):
    # First produce a complete 3-angle file
    engine, trace, plan = build(tmp_path, num_angles=3)
    engine.run()
    path = tmp_path / "out.sras"
    original = path.read_bytes()

    sras = SrasFile(path)
    statuses = sras.angle_status()
    target = statuses[1]
    resume = ResumeState(
        path=path,
        targets=[ResumeTarget(target.index, target.bg_offset,
                              target.data_offset, target.n_rows,
                              target.angle_deg)],
        samples_per_frame=SPF,
    )

    engine2, trace2, _ = build(tmp_path, num_angles=3, resume=resume)
    result = engine2.run()

    assert result.angles_acquired == [1]
    # Only the middle angle's rows were re-acquired
    assert result.rows_written == plan.per_angle[1].n_rows
    rewritten = path.read_bytes()
    assert len(rewritten) == len(original)
    # Angle 0's block is untouched; angle 1's changed — background included,
    # since a re-acquired angle captures a fresh one over the old.
    a1_start = target.bg_offset
    a1_end = target.data_offset + target.row_bytes * target.n_rows
    assert rewritten[:a1_start] == original[:a1_start]
    assert rewritten[a1_start:a1_end] != original[a1_start:a1_end]
    assert rewritten[a1_end:] == original[a1_end:]

    # The new background is angle 1's own, written in place of its old one.
    reread = SrasFile(path)
    assert reread.backgrounds[1] == background_record(0, SPF)
    assert reread.backgrounds[0] == SrasFile(path).backgrounds[0]
    assert [s.status for s in reread.angle_status()] == ["OK"] * 3


def test_resume_rejects_a_background_that_would_shift_the_file(tmp_path):
    """A re-acquired angle's background must fit the room the file has.

    Nothing else in the file records where an angle's rows begin, so a longer
    or shorter background would push every row behind it out of position.
    """
    engine, _, _ = build(tmp_path, num_angles=2)
    engine.run()
    path = tmp_path / "out.sras"
    target = SrasFile(path).angle_status()[0]
    resume = ResumeState(
        path=path,
        targets=[ResumeTarget(0, target.bg_offset, target.data_offset,
                              target.n_rows, target.angle_deg)],
        samples_per_frame=SPF,
    )
    engine2, _, _ = build(tmp_path, num_angles=2, resume=resume)
    # A scope that hands back a longer record than the file was written with
    engine2._scope.transfer_curve = lambda: bytes(SPF + 4)

    with pytest.raises(RuntimeError, match="shift every row"):
        engine2.run()


def test_resume_record_length_mismatch_rejected(tmp_path):
    engine, trace, plan = build(tmp_path)
    engine.run()
    path = tmp_path / "out.sras"
    resume = ResumeState(path=path,
                         targets=[ResumeTarget(0, 0, 0, 1, 0.0)],
                         samples_per_frame=SPF + 1)   # scope changed
    engine2, _, _ = build(tmp_path, resume=resume)
    with pytest.raises(RuntimeError, match="record length"):
        engine2.run()


# ── Burst acquisition ────────────────────────────────────────────────────────

# 6 rows × 4 frames/row; max_frames=14 gives 14//4 = 3 rows per burst, so the
# angle needs two bursts and the second is not a whole burst wide.
BURST_PLAN = dict(y_delta=0.05)
BURST_MAX_FRAMES = 14


def test_burst_and_serial_produce_identical_files(tmp_path):
    """The whole point: burst mode must be a pure acquisition optimisation."""
    plan = make_plan(**BURST_PLAN)
    assert plan.per_angle[0].n_rows == 6 and plan.per_angle[0].n_frames == 4

    serial, _, _ = build(tmp_path, plan=plan, out_name="serial.sras")
    serial.run()
    burst, _, _ = build(tmp_path, plan=plan, out_name="burst.sras",
                        burst_mode=True, max_frames=BURST_MAX_FRAMES)
    burst.run()

    assert (tmp_path / "burst.sras").read_bytes() == \
           (tmp_path / "serial.sras").read_bytes()


def test_burst_multi_angle_file_is_complete(tmp_path):
    plan = make_plan(num_angles=3, **BURST_PLAN)
    engine, trace, _ = build(tmp_path, plan=plan, burst_mode=True,
                             max_frames=BURST_MAX_FRAMES)
    result = engine.run()

    assert result.rows_written == plan.total_rows
    assert result.angles_acquired == [0, 1, 2]
    sras = SrasFile(tmp_path / "out.sras")
    assert [s.status for s in sras.angle_status()] == ["OK"] * 3


def test_burst_gates_the_flyback_and_runs_once_per_burst(tmp_path):
    plan = make_plan(**BURST_PLAN)
    engine, trace, _ = build(tmp_path, plan=plan, burst_mode=True,
                             max_frames=BURST_MAX_FRAMES)
    engine.run()

    # Two bursts (3 + 3 rows), two preflight acquisitions, one background.
    runs = [c for c in trace.of("write") if c[1] == "ACQuire:STATE RUN"]
    assert len(runs) == 5

    # Every acquiring pass is bracketed by an arm/disarm, so the gate is low
    # for each flyback.  6 rows + 1 preflight pass = 7 arms.
    gate = [c[2] for c in trace.of("arm_scan_gate")]
    assert gate.count(True) == 7
    # No two arms without a disarm between them — that is what would let a
    # flyback into the acquisition.  (A repeated disarm is just defensive.)
    for a, b in zip(gate, gate[1:], strict=False):
        assert not (a and b), f"acquiring pass with no disarm before it: {gate}"
    assert gate[-1] is False, "scan left the gate armed"

    # One bulk transfer per data channel per burst, none per row.
    assert [(c[1], c[2]) for c in trace.of("transfer_fastframe_bulk")] == [
        (1, 12), (4, 12), (1, 12), (4, 12)]
    assert trace.count("transfer_fastframe") == 0


def test_burst_preflight_rejects_a_leaky_gate(tmp_path):
    plan = make_plan(**BURST_PLAN)
    engine, trace, _ = build(tmp_path, plan=plan, burst_mode=True,
                             max_frames=BURST_MAX_FRAMES)
    stage = engine._stage

    # A gate that ignores the disable request — the failure mode the preflight
    # exists to catch (TRIGOUT_GATE_OFF set to the wrong mode value).
    def stuck_gate(axis, armed, verify=True):
        trace.record("arm_scan_gate", axis, bool(armed))
        stage.gate_armed = True
    stage.arm_scan_gate = stuck_gate

    with pytest.raises(RuntimeError, match="not idling low"):
        engine.run()


def test_burst_preflight_rejects_a_dark_laser(tmp_path):
    """A gate that never fires would let a leak check pass vacuously."""
    plan = make_plan(**BURST_PLAN)
    engine, trace, _ = build(tmp_path, plan=plan, burst_mode=True,
                             max_frames=BURST_MAX_FRAMES)
    engine._stage.attach_scope(None)      # no pulses ever reach the scope

    with pytest.raises(RuntimeError, match="no frames acquired"):
        engine.run()


def _clip_one_row(engine, which_pass=2, lost=1):
    """Make one acquiring pass come up `lost` frames short.

    `which_pass` counts acquiring passes from 1, so the default clips the
    second data row (row 2) — far enough in that a mishandled short row shows
    up as a shift in the rows after it.
    """
    scope = engine._scope
    real_acquire = scope.acquire_frames
    passes = {"n": 0}

    def clipped(n):
        if scope._running:
            passes["n"] += 1
            if passes["n"] == which_pass:
                n -= lost
        real_acquire(n)
    scope.acquire_frames = clipped


@pytest.mark.parametrize("burst_mode", [False, True])
def test_short_row_is_padded_to_declared_frame_count(tmp_path, burst_mode):
    """A clipped row must not shift every later row in the file.

    v6 declares n_frames per row up front and has no per-row length, so an
    under-triggered row has to be squared up.  In burst mode this also proves
    the splitter advances by what actually arrived, not by n_frames.
    """
    warnings = []
    plan = make_plan(**BURST_PLAN)
    engine, trace, _ = build(
        tmp_path, plan=plan, burst_mode=burst_mode,
        max_frames=BURST_MAX_FRAMES,
        callbacks=ScanCallbacks(on_status=warnings.append))
    # The preflight is covered by its own tests; skipping it keeps the
    # acquiring-pass count below identical in both modes.
    engine._preflight_done = True

    _clip_one_row(engine)

    result = engine.run()

    assert result.rows_written == 6
    assert any("Row 2: 3 frames acquired, 4 expected" in w for w in warnings)
    assert any("zero-padded" in w for w in warnings)
    sras = SrasFile(tmp_path / "out.sras")
    assert [s.status for s in sras.angle_status()] == ["OK"]
    # The padding lands at the end of the short row, not in the next one.
    assert bytes(sras.load_row(0, 1, 0)[-1]) == bytes(SPF)
    assert bytes(sras.load_row(0, 2, 0)[0]) != bytes(SPF)


# ── Strict row packing ───────────────────────────────────────────────────────

@pytest.mark.parametrize("burst_mode", [False, True])
def test_strict_row_packing_aborts_on_a_short_row(tmp_path, burst_mode):
    """Strict mode fails the scan instead of silently squaring a row up.

    The default padding keeps the file readable but makes a mis-triggered row
    indistinguishable from a good one after the fact, since v6 records no
    per-row frame count. Strict mode trades the salvaged rows for knowing.
    """
    plan = make_plan(**BURST_PLAN)
    engine, _, _ = build(tmp_path, plan=plan, burst_mode=burst_mode,
                         max_frames=BURST_MAX_FRAMES, strict_rows=True)
    engine._preflight_done = True
    _clip_one_row(engine)

    with pytest.raises(RuntimeError, match="Row 2: 3 frames acquired, 4 expected"):
        engine.run()


@pytest.mark.parametrize("burst_mode", [False, True])
def test_strict_row_packing_does_not_disturb_a_clean_scan(tmp_path, burst_mode):
    """Strict mode is inert when every row acquires what it declared."""
    plan = make_plan(**BURST_PLAN)
    engine, _, _ = build(tmp_path, plan=plan, burst_mode=burst_mode,
                         max_frames=BURST_MAX_FRAMES, strict_rows=True)
    engine._preflight_done = True

    result = engine.run()

    assert result.rows_written == plan.total_rows
    sras = SrasFile(tmp_path / "out.sras")
    assert [s.status for s in sras.angle_status()] == ["OK"]


def test_strict_row_packing_writes_nothing_for_the_failed_row(tmp_path):
    """The abort must not leave a half-written row behind.

    CH1 leads SCAN_CHANNELS, so the frame count is known before any of the
    row's channels are written — the file should end on a whole-row boundary.
    """
    plan = make_plan(**BURST_PLAN)
    engine, _, _ = build(tmp_path, plan=plan, strict_rows=True)
    engine._preflight_done = True
    _clip_one_row(engine)

    with pytest.raises(RuntimeError, match="Strict row packing"):
        engine.run()

    # Row 1 was written in full; row 2 aborted before writing anything, so
    # the file ends exactly on a row boundary.
    sras = SrasFile(tmp_path / "out.sras")
    written = (tmp_path / "out.sras").stat().st_size - sras.angle_data_offset(0)
    assert written == sras.row_bytes(0)


def test_engine_imports_without_qt():
    """The engine must be usable from a non-Qt front end."""
    import subprocess
    import sys
    code = (
        "import sys;"
        "sys.modules['PyQt6'] = None;"
        "import core.scan_engine, core.rotation, core.scope_sras,"
        " core.scan_resume, core.sras_format, core.scan_geometry;"
        "print('ok')"
    )
    out = subprocess.run([sys.executable, "-c", code], capture_output=True,
                         text=True, cwd=str(__import__('pathlib').Path(__file__).parent.parent))
    assert out.returncode == 0, out.stderr
    assert "ok" in out.stdout