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.
"""
from __future__ import annotations


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)):
        self._t = trace
        self.am_homed = list(homed)
        self.am_enabled = list(enabled)
        self.positions = [0.0, 0.0]

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

    def home_axis(self, axis, timeout=0.0):
        self._t.record("home_axis", axis)
        self.am_homed[0 if axis == 0x21 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)

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


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, n_frames=4):
        self._t = trace
        self.samples_per_frame = samples_per_frame
        self._n_frames = n_frames
        self._acq_polls = 0
        self.frame_seq = 0

    # -- writes / queries ---------------------------------------------------
    def write(self, cmd):
        self._t.record("write", cmd)

    def query(self, cmd):
        self._t.record("query", cmd)
        if cmd == "ACQuire:STATE?":
            self._acq_polls += 1
            return "0"                      # background average finished
        if cmd == "ACQuire:NUMFRAMESACQuired?":
            return str(self._n_frames)
        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)
        self._n_frames = n

    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 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")
        return bytes(range(self.samples_per_frame))

    def transfer_fastframe(self, parse=True):
        self._t.record("transfer_fastframe", self._source)
        frames = []
        for i in range(self._n_frames):
            frames.append(bytes((self.frame_seq + i + s) % 256
                                for s in range(self.samples_per_frame)))
        self.frame_seq += 1
        return frames

    # 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