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 SCAN_CHANNELS, SrasFile
from fakes import FakeScope, FakeStage, FakeT3R, Trace

SPF = 8


def make_plan(num_angles=1):
    # Small ROI well inside the stage limits: 1 row, few frames per angle.
    return build_plan(40.0, 30.0, 0.02, 0.005, 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, **kw):
    trace = Trace()
    stage = FakeStage(trace)
    scope = FakeScope(trace, samples_per_frame=SPF)
    t3r = FakeT3R(trace, **kw)
    rotator = RotationAxis(t3r, RotationSettings())
    plan = make_plan(num_angles)
    engine = ScanEngine(stage, scope, rotator, plan, tmp_path / "out.sras",
                        resume=resume,
                        callbacks=callbacks or ScanCallbacks())
    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.background == bytes(range(SPF))


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()
    stage = FakeStage(trace)
    scope = FakeScope(trace, samples_per_frame=SPF)
    # 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.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 (fresh frame data)
    a1_start, a1_end = target.data_offset, 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:]


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, 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()


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