scanengine-3/tests/fakes.py

"""Recording fake hardware for headless ScanEngine tests.

Each fake records an ordered call trace, so a test can assert the exact
command sequence the engine issues — the property that matters when the
real rig isn't available.

The stage and scope are wired together the way the rig is: an X move at scan
velocity with the trigger gate armed feeds frames into a running acquisition,
at the real 20 kHz / 100 mm/s rate.  Per-row and burst acquisition therefore
get their frame counts from the same model, which is what makes a
byte-identity comparison between the two paths meaningful — and it means a
gate the engine forgets to drop shows up as extra frames instead of passing
silently.
"""
from __future__ import annotations

from core.scan_engine import (
    AXIS_X, LASER_FREQ_HZ, SCAN_RAMP_BUFFER_MM, SCAN_RAMP_MM,
    SCAN_VELOCITY_MM_S,
)

RAMP_TOTAL_MM = SCAN_RAMP_MM + SCAN_RAMP_BUFFER_MM


class Trace:
    """Ordered record of hardware calls, shared by all fakes in one test."""

    def __init__(self):
        self.calls: list[tuple] = []

    def record(self, *entry):
        self.calls.append(entry)

    def names(self) -> list[str]:
        return [c[0] for c in self.calls]

    def of(self, name: str) -> list[tuple]:
        return [c for c in self.calls if c[0] == name]

    def count(self, name: str) -> int:
        return len(self.of(name))


class FakeStage:
    """Stands in for ThorlabsServoDriver."""

    def __init__(self, trace: Trace, homed=(True, True), enabled=(True, True),
                 scope=None):
        self._t = trace
        self.am_homed = list(homed)
        self.am_enabled = list(enabled)
        self.positions = [0.0, 0.0]
        self._scope = scope
        self.gate_armed = False

    def attach_scope(self, scope):
        """Route gated motion into `scope`, as the TRIGOUT pin does on the rig."""
        self._scope = scope

    def enable_axis(self, axis):
        self._t.record("enable_axis", axis)
        self.am_enabled[0 if axis == AXIS_X else 1] = True

    def home_axis(self, axis, timeout=0.0):
        self._t.record("home_axis", axis)
        self.am_homed[0 if axis == AXIS_X else 1] = True

    def set_velocity_params(self, axis, max_velocity=None, acceleration=None):
        self._t.record("set_velocity_params", axis, max_velocity, acceleration)

    def set_trigger_trigout_maxv(self, axis):
        self._t.record("set_trigger_trigout_maxv", axis)
        if axis == AXIS_X:
            self.gate_armed = True

    def set_trigger_gate_off(self, axis):
        self._t.record("set_trigger_gate_off", axis)
        if axis == AXIS_X:
            self.gate_armed = False

    def arm_scan_gate(self, axis, armed, verify=True):
        self._t.record("arm_scan_gate", axis, bool(armed))
        if axis == AXIS_X:
            self.gate_armed = bool(armed)

    def move_axis_absolute(self, axis, pos, timeout=0.0):
        idx = 0 if axis == AXIS_X else 1
        prev = self.positions[idx]
        self._t.record("move_axis_absolute", axis, round(pos, 6))
        self.positions[idx] = pos

        # The gate is high only at max velocity, i.e. over the move minus its
        # two ramps — direction-agnostic, so a flyback the engine failed to
        # gate off produces frames instead of quietly producing none.
        if axis == AXIS_X and self.gate_armed and self._scope is not None:
            at_speed_mm = abs(pos - prev) - 2 * RAMP_TOTAL_MM
            if at_speed_mm > 0:
                self._scope.acquire_frames(
                    round(at_speed_mm * LASER_FREQ_HZ / SCAN_VELOCITY_MM_S))


def background_record(n: int, samples_per_frame: int) -> bytes:
    """The waveform FakeScope returns from its n-th background capture.

    Every angle captures its own, so the tests need to tell one from the
    next: a background that landed under the wrong angle would otherwise
    look exactly like the right one.
    """
    return bytes((n * 17 + s) % 256 for s in range(samples_per_frame))


class FakeScope:
    """Stands in for TektronixOscilloscopeBase.

    Returns deterministic frame bytes so the written file can be compared
    against an expected byte pattern.
    """

    def __init__(self, trace: Trace, samples_per_frame=8, max_frames=4096):
        self._t = trace
        self.samples_per_frame = samples_per_frame
        self.max_frames = max_frames
        self._acq_polls = 0
        self._running = False
        self._acquired = 0
        self._backgrounds_taken = 0
        # Per-channel running frame index.  Frame content is a function of
        # (channel, index) alone, so the same total frame sequence yields the
        # same bytes however it is chopped into transfers.
        self._next_frame: dict[int, int] = {}

    # -- driven by FakeStage ------------------------------------------------
    def acquire_frames(self, n):
        if self._running:
            self._acquired += n

    # -- writes / queries ---------------------------------------------------
    def write(self, cmd):
        self._t.record("write", cmd)
        if cmd == "ACQuire:STATE RUN":
            self._running = True
            self._acquired = 0
        elif cmd == "ACQuire:STATE STOP":
            self._running = False

    def query(self, cmd):
        self._t.record("query", cmd)
        if cmd == "ACQuire:STATE?":
            self._acq_polls += 1
            # STOPAfter SEQuence self-stops when the sequence completes, so
            # reporting "stopped" and staying armed would be inconsistent.
            self._running = False
            return "0"                      # background average finished
        if cmd == "ACQuire:NUMFRAMESACQuired?":
            return str(self._acquired)
        return ""

    # -- typed setters used by core.scope_sras ------------------------------
    def set_trigger_source(self, ch):
        self._t.record("set_trigger_source", ch)

    def set_trigger_slope(self, slope):
        self._t.record("set_trigger_slope", slope)

    def set_trigger_level(self, ch, level):
        self._t.record("set_trigger_level", ch, level)

    def set_trigger_mode(self, mode):
        self._t.record("set_trigger_mode", mode)

    def set_acquire_mode(self, mode):
        self._t.record("set_acquire_mode", mode)

    def set_fastframe_state(self, on):
        self._t.record("set_fastframe_state", on)

    def set_fastframe_count(self, n):
        self._t.record("set_fastframe_count", n)

    def get_fastframe_state(self):
        return 1

    def get_fastframe_max_frames(self):
        self._t.record("get_fastframe_max_frames")
        return self.max_frames

    def set_sample_rate(self, sr):
        self._t.record("set_sample_rate", sr)

    def get_record_length(self):
        return self.samples_per_frame

    def set_data_source(self, ch):
        self._t.record("set_data_source", ch)
        self._source = ch

    def set_data_encoding(self, encoding):
        self._t.record("set_data_encoding", encoding)

    def set_data_width(self, width):
        self._t.record("set_data_width", width)

    def query_wfmoutpre(self):
        return f"WFMOUTPRE:CH{self._source};YMULT 1.5625E-3;YOFF -87.04;YZERO 0.0"

    def transfer_curve(self):
        self._t.record("transfer_curve")
        n = self._backgrounds_taken
        self._backgrounds_taken += 1
        return background_record(n, self.samples_per_frame)

    def _frames(self, ch, count):
        spf = self.samples_per_frame
        start = self._next_frame.get(ch, 0)
        self._next_frame[ch] = start + count
        return [bytes((ch * 31 + g + s) % 256 for s in range(spf))
                for g in range(start, start + count)]

    def transfer_fastframe(self, parse=True, byte_count=1, signed=True,
                           byte_order='MSB'):
        self._t.record("transfer_fastframe", self._source)
        return self._frames(self._source, self._acquired)

    def transfer_fastframe_bulk(self, frame_count, samples_per_frame,
                                bytes_per_sample=1):
        self._t.record("transfer_fastframe_bulk", self._source, frame_count)
        return bytearray(b"".join(self._frames(self._source, frame_count)))

    # channel config (only used by configure_channels)
    def set_channel_label_name(self, ch, name):
        self._t.record("set_channel_label_name", ch, name)

    def set_channel_scale(self, ch, v):
        self._t.record("set_channel_scale", ch, v)

    def set_channel_position(self, ch, v):
        self._t.record("set_channel_position", ch, v)

    def set_channel_termination(self, ch, v):
        self._t.record("set_channel_termination", ch, v)

    def set_channel_coupling(self, ch, v):
        self._t.record("set_channel_coupling", ch, v)

    def set_channel_bandwidth(self, ch, v):
        self._t.record("set_channel_bandwidth", ch, v)


class FakeT3R:
    """Stands in for the (Qt-free) T3RDriver, for RotationAxis."""

    GR_AXIS_CH = 3
    MOTOR_FULL_STEPS_PER_REV = 200
    GEAR_TEETH_MOTOR = 10
    GEAR_TEETH_STAGE = 125

    def __init__(self, trace: Trace, is_open=True, motion_completes=True):
        self._t = trace
        self.is_open = is_open
        self._motion_completes = motion_completes

    def set_microstep(self, ch, micro):
        self._t.record("t3r_set_microstep", ch, micro)

    def set_current(self, ch, run_ma, hold_ma, ihold):
        self._t.record("t3r_set_current", ch, run_ma, hold_ma, ihold)

    def enable(self, ch):
        self._t.record("t3r_enable", ch)

    def steps_for_angle(self, angle_deg, microsteps):
        ratio = self.GEAR_TEETH_STAGE / self.GEAR_TEETH_MOTOR
        return round(self.MOTOR_FULL_STEPS_PER_REV * microsteps * ratio
                     * angle_deg / 360.0)

    def rotate_stage(self, angle_deg, microsteps, velocity, accel):
        self._t.record("t3r_rotate", round(angle_deg, 6))

    def wait_motion_done(self, ch, timeout):
        self._t.record("t3r_wait_motion_done", ch)
        return self._motion_completes