csaxs_bec/tests/tests_devices/test_panda.py

"""Module for testing the PandaBoxCSAXS and PandaBoxOMNY devices."""

# pylint: skip-file
from __future__ import annotations

import threading
from unittest import mock

import numpy as np
import pytest
from ophyd import Staged
from ophyd_devices.devices.panda_box.panda_box import PandaState
from ophyd_devices.interfaces.base_classes.psi_device_base import DeviceStoppedError
from pandablocks.responses import FrameData

from csaxs_bec.devices.panda_box.panda_box import PandaBoxCSAXS
from csaxs_bec.devices.panda_box.panda_box_omny import PandaBoxOMNY
from csaxs_bec.devices.utils.utils import fetch_scan_info


def _make_frame_data(n_rows: int, start: int = 0) -> FrameData:
    """Build a synthetic FrameData batch with `n_rows` rows for the frame_counter /
    gate_detector_active signals used by the raw-stream tests below."""
    data = np.array(
        [(np.float64(i), np.float64(i % 2)) for i in range(start, start + n_rows)],
        dtype=[("COUNTER1.OUT.Value", "<f8"), ("PCAP.BITS0.Value", "<f8")],
    )
    return FrameData(data)


@pytest.fixture
def panda_omny():
    dev_name = "panda_omny"
    dev = PandaBoxOMNY(
        name=dev_name,
        host="omny-panda-box.psi.ch",
        signal_alias={
            "FMC_IN.VAL1.Min": "cap_voltage_fzp_y_min",
            "FMC_IN.VAL1.Max": "cap_voltage_fzp_y_max",
            "FMC_IN.VAL1.Mean": "cap_voltage_fzp_y_mean",
            "FMC_IN.VAL2.Min": "cap_voltage_fzp_x_min",
            "FMC_IN.VAL2.Max": "cap_voltage_fzp_x_max",
            "FMC_IN.VAL2.Mean": "cap_voltage_fzp_x_mean",
        },
    )
    yield dev


@pytest.fixture
def panda_omny_raw_stream():
    dev_name = "panda_omny_raw"
    dev = PandaBoxOMNY(
        name=dev_name,
        host="omny-panda-box.psi.ch",
        signal_alias={
            "PCAP.BITS0.Value": "gate_detector_active",
            "COUNTER1.OUT.Value": "frame_counter",
        },
        raw_stream_mode=True,
        raw_stream_flush_row_count=5,
        # Effectively disable the time-based flush so tests are deterministic and only
        # exercise the row-count threshold.
        raw_stream_flush_interval=1000,
    )
    yield dev


@pytest.fixture
def panda_csaxs():
    dev_name = "panda_csaxs"
    dev = PandaBoxCSAXS(name=dev_name, host="csaxs-panda-box.psi.ch")
    yield dev


def test_panda_omny(panda_omny):
    assert panda_omny.name == "panda_omny"
    assert panda_omny.host == "omny-panda-box.psi.ch"

    all_signal_names = [name for name, _ in panda_omny.data.signals]
    # Check that the signal aliases are correctly set up
    assert "cap_voltage_fzp_y_min" in all_signal_names
    assert "cap_voltage_fzp_y_max" in all_signal_names
    assert "cap_voltage_fzp_y_mean" in all_signal_names
    assert "cap_voltage_fzp_x_min" in all_signal_names
    assert "cap_voltage_fzp_x_max" in all_signal_names
    assert "cap_voltage_fzp_x_mean" in all_signal_names

    # Check that the original signal names are not present
    assert "FMC_IN.VAL1.Min" not in all_signal_names
    assert "FMC_IN.VAL1.Max" not in all_signal_names
    assert "FMC_IN.VAL1.Mean" not in all_signal_names
    assert "FMC_IN.VAL2.Min" not in all_signal_names
    assert "FMC_IN.VAL2.Max" not in all_signal_names
    assert "FMC_IN.VAL2.Mean" not in all_signal_names

    assert panda_omny._acquisition_group == "burst"
    assert panda_omny._timeout_on_completed == 10


@pytest.mark.parametrize(
    "scan_type, frames_per_trigger, expected_acquisition_group",
    [
        ("hardware_triggered", 1, "fly"),
        ("hardware_triggered", 5, "fly"),
        ("software_triggered", 10, "burst"),
        ("software_triggered", 1, "monitored"),  # Default case
    ],
)
def test_panda_omny_stage(panda_omny, scan_type, frames_per_trigger, expected_acquisition_group):
    # Check that the stage signal is present and has the correct PV
    assert len(panda_omny._status_callbacks) == 0

    panda_omny.scan_info.msg.info["scan_type"] = scan_type
    panda_omny.scan_info.msg.info["frames_per_trigger"] = frames_per_trigger
    panda_omny.stage()

    assert panda_omny._acquisition_group == expected_acquisition_group
    assert panda_omny.staged == Staged.yes


def test_panda_omny_complete(panda_omny):
    """Test the on_complete method of the PandaBoxCSAXS device."""
    panda_omny.scan_info.msg.info["num_points"] = 1
    panda_omny.scan_info.msg.info["frames_per_trigger"] = 1
    panda_omny.scan_parameters = fetch_scan_info(panda_omny.scan_info)

    panda_omny._timeout_on_completed = 0.5  # Set a short timeout for testing

    def _mock_return_captured(*args, **kwargs):
        return ["=0"]

    # Timeout Error on complete
    with (
        mock.patch.object(panda_omny, "send_raw", side_effect=_mock_return_captured),
        mock.patch.object(panda_omny, "_disarm", return_value=None) as mock_disarm,
    ):
        status = panda_omny.on_complete()
        assert status.done is False
        assert status.success is False

        with pytest.raises(TimeoutError):
            status.wait(timeout=4)
        mock_disarm.assert_called_once()

    # Successful complete
    panda_omny._timeout_on_completed = 5
    with (
        mock.patch.object(panda_omny, "send_raw", side_effect=[["=0"], ["=0"], ["=1"]]),
        mock.patch.object(panda_omny, "_disarm", return_value=None) as mock_disarm,
    ):
        status = panda_omny.on_complete()
        assert status.done is False
        assert status.success is False

        status.wait(timeout=4)
        mock_disarm.assert_called_once()
        assert status.done is True
        assert status.success is True


def test_panda_omny_default_raw_stream_mode_is_off(panda_omny):
    """New devices must default to legacy behavior unless raw_stream_mode is explicitly set."""
    assert panda_omny.raw_stream_mode is False


def test_panda_omny_legacy_receive_frame_data_is_unbuffered(panda_omny):
    """Backward-compatibility regression: without raw_stream_mode, each FrameData must still
    be published immediately, exactly as before this feature was added."""
    with mock.patch.object(panda_omny.data, "put") as mock_put:
        for i in range(3):
            panda_omny._receive_frame_data(_make_frame_data(2, start=i * 2))
        assert mock_put.call_count == 3


def test_panda_omny_raw_stream_acquisition_group(panda_omny_raw_stream):
    """raw_stream_mode must select the free_running acquisition group regardless of
    scan_type/frames_per_trigger."""
    panda_omny_raw_stream.scan_info.msg.info["scan_type"] = "software_triggered"
    panda_omny_raw_stream.scan_info.msg.info["frames_per_trigger"] = 10
    panda_omny_raw_stream.stage()

    assert panda_omny_raw_stream._acquisition_group == "free_running"
    assert panda_omny_raw_stream.staged == Staged.yes


def test_panda_omny_raw_stream_coalesces_by_row_count(panda_omny_raw_stream):
    """Multiple small FrameData batches must be coalesced into a single .data.put() call
    once the row-count threshold is reached, not published one-by-one."""
    dev = panda_omny_raw_stream
    with mock.patch.object(dev.data, "put") as mock_put:
        dev._receive_frame_data(_make_frame_data(2, start=0))
        dev._receive_frame_data(_make_frame_data(2, start=2))
        mock_put.assert_not_called()  # 4 rows buffered so far, threshold is 5

        dev._receive_frame_data(_make_frame_data(2, start=4))
        mock_put.assert_called_once()  # 6 rows >= 5, flush triggered

        out = mock_put.call_args.args[0]
        assert [v for v in out["frame_counter"]["value"]] == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
        assert [v for v in out["gate_detector_active"]["value"]] == [0.0, 1.0, 0.0, 1.0, 0.0, 1.0]


def test_panda_omny_raw_stream_gate_bit_is_masked_from_shared_word(panda_omny_raw_stream):
    """PCAP.BITS0 is a 32-bit word shared with other captured bit_out signals; only bit 0
    (TTLIN1.VAL, confirmed via CAPTURE_WORD?/OFFSET? against the omny-panda hardware) is the
    detector-trigger gate. gate_detector_active must be masked down to that bit rather than
    passed through as the raw word."""
    dev = panda_omny_raw_stream
    data = np.array(
        # bit 0 set (gate active) plus other bits from something else sharing BITS0
        [(0.0, 5.0), (1.0, 4.0), (2.0, 3.0)],
        dtype=[("COUNTER1.OUT.Value", "<f8"), ("PCAP.BITS0.Value", "<f8")],
    )
    with mock.patch.object(dev.data, "put") as mock_put:
        dev._receive_frame_data(FrameData(data))
        dev._flush_raw_stream_buffer()

        out = mock_put.call_args.args[0]
        assert [v for v in out["gate_detector_active"]["value"]] == [1, 0, 1]


def test_panda_omny_raw_stream_flush_on_demand(panda_omny_raw_stream):
    """Buffered data below the threshold must still be flushed on demand (as on_complete/
    on_unstage do), and flushing an empty buffer must not publish again."""
    dev = panda_omny_raw_stream
    with mock.patch.object(dev.data, "put") as mock_put:
        dev._receive_frame_data(_make_frame_data(2, start=0))
        mock_put.assert_not_called()

        dev._flush_raw_stream_buffer()
        mock_put.assert_called_once()
        out = mock_put.call_args.args[0]
        assert len(out["frame_counter"]["value"]) == 2

        dev._flush_raw_stream_buffer()
        mock_put.assert_called_once()  # still 1: nothing left to flush


def test_panda_omny_unstage_flushes_pending_raw_stream_data(panda_omny_raw_stream):
    dev = panda_omny_raw_stream
    dev.stage()
    with (
        mock.patch.object(dev.data, "put") as mock_put,
        mock.patch.object(dev, "_disarm", return_value=None),
    ):
        dev._receive_frame_data(_make_frame_data(2, start=0))
        mock_put.assert_not_called()

        dev.unstage()
        mock_put.assert_called_once()


def test_panda_omny_complete_free_running_waits_for_frame_counter(panda_omny_raw_stream):
    """Completion in free-running mode waits for COUNTER1 (frame_counter) to reach the exact
    expected frame count -- PCAP.CAPTURED can't be used since it free-runs continuously off
    CLOCK1 and never naturally stops increasing while armed."""
    dev = panda_omny_raw_stream
    dev.scan_info.msg.info["num_points"] = 1
    dev.scan_info.msg.info["frames_per_trigger"] = 1
    dev.scan_parameters = fetch_scan_info(dev.scan_info)
    dev._timeout_on_completed = 5

    with (
        mock.patch.object(
            dev, "send_raw", side_effect=[["=0"], ["=0"], ["=1"]]
        ) as mock_send_raw,
        mock.patch.object(dev, "_disarm", return_value=None) as mock_disarm,
        mock.patch.object(dev, "_flush_raw_stream_buffer") as mock_flush,
    ):
        status = dev.on_complete()
        status.wait(timeout=4)
        assert status.done is True
        assert status.success is True
        mock_disarm.assert_called_once()
        mock_flush.assert_called_once()
        assert mock_send_raw.call_args.args[0] == "COUNTER1.OUT?"


def test_panda_omny_complete_free_running_timeout(panda_omny_raw_stream):
    """If the frame count never reaches the expected total, the overall timeout must still
    fire."""
    dev = panda_omny_raw_stream
    dev.scan_info.msg.info["num_points"] = 1
    dev.scan_info.msg.info["frames_per_trigger"] = 1
    dev.scan_parameters = fetch_scan_info(dev.scan_info)
    dev._timeout_on_completed = 0.3

    with (
        mock.patch.object(dev, "send_raw", side_effect=lambda *a, **k: ["=0"]),
        mock.patch.object(dev, "_disarm", return_value=None) as mock_disarm,
        mock.patch.object(dev, "_flush_raw_stream_buffer") as mock_flush,
    ):
        status = dev.on_complete()
        with pytest.raises(TimeoutError):
            status.wait(timeout=4)
        mock_disarm.assert_called_once()
        mock_flush.assert_called_once()


def _deliver_ready_after(dev, delay: float) -> threading.Timer:
    """Simulate the PandA data thread delivering the READY event `delay` seconds from now."""
    timer = threading.Timer(delay, dev._run_status_callbacks, args=(PandaState.READY,))
    timer.daemon = True
    timer.start()
    return timer


def test_panda_omny_pre_scan_arms_when_ready(panda_omny, capsys):
    dev = panda_omny
    dev.pre_scan_timeout = 2
    with (
        mock.patch.object(dev, "_arm") as mock_arm,
        mock.patch.object(dev, "_reset_panda") as mock_reset,
    ):
        _deliver_ready_after(dev, 0.05)
        status = dev.on_pre_scan()
        status.wait(timeout=4)
        assert status.success is True
        mock_arm.assert_called_once()
        mock_reset.assert_not_called()
    assert "Forcing abort/reset" not in capsys.readouterr().out


def test_panda_omny_pre_scan_retries_once_then_succeeds(panda_omny, capsys):
    """No READY on the first attempt: the box is force-reset and arming retried. The retry is
    printed and logged."""
    dev = panda_omny
    dev.pre_scan_timeout = 0.2

    def _reset_then_ready():
        _deliver_ready_after(dev, 0.05)

    with (
        mock.patch.object(dev, "_arm") as mock_arm,
        mock.patch.object(dev, "_reset_panda", side_effect=_reset_then_ready) as mock_reset,
        mock.patch("csaxs_bec.devices.panda_box.panda_box_omny.logger") as mock_logger,
    ):
        status = dev.on_pre_scan()
        status.wait(timeout=4)
        assert status.success is True
        mock_reset.assert_called_once()
        mock_arm.assert_called_once()
        assert any("Forcing abort/reset" in str(c) for c in mock_logger.warning.call_args_list)
        assert any("armed successfully on attempt 2" in str(c) for c in mock_logger.info.call_args_list)
    assert "Forcing abort/reset" in capsys.readouterr().out


def test_panda_omny_pre_scan_raises_after_retry_fails(panda_omny, capsys):
    dev = panda_omny
    dev.pre_scan_timeout = 0.1
    with (
        mock.patch.object(dev, "_arm") as mock_arm,
        mock.patch.object(dev, "_reset_panda") as mock_reset,
        mock.patch("csaxs_bec.devices.panda_box.panda_box_omny.logger") as mock_logger,
    ):
        status = dev.on_pre_scan()
        with pytest.raises(RuntimeError, match="could not be armed after 2 attempts"):
            status.wait(timeout=4)
        mock_arm.assert_not_called()
        # once for the forced reset before the retry, once to leave the box clean at the end
        assert mock_reset.call_count == 2
        mock_logger.error.assert_called_once()
    assert capsys.readouterr().out.count("Forcing abort/reset") == 1
    # stale status callbacks of the failed attempts are removed again
    assert dev._status_callbacks == {}


def test_panda_omny_pre_scan_stop_is_not_retried(panda_omny, capsys):
    dev = panda_omny
    dev.pre_scan_timeout = 0.3
    with (
        mock.patch.object(dev, "_arm") as mock_arm,
        mock.patch.object(dev, "_reset_panda") as mock_reset,
    ):
        status = dev.on_pre_scan()
        threading.Event().wait(0.05)
        dev._stop_stoppable_status_objects()
        with pytest.raises(DeviceStoppedError):
            status.wait(timeout=4)
        # give a wrongly retrying task thread the time to show up
        threading.Event().wait(0.5)
        mock_arm.assert_not_called()
        mock_reset.assert_not_called()
    assert "Forcing abort/reset" not in capsys.readouterr().out


def test_panda_omny_raw_stream_signal_alias(panda_omny_raw_stream):
    all_signal_names = [name for name, _ in panda_omny_raw_stream.data.signals]
    assert "gate_detector_active" in all_signal_names
    assert "frame_counter" in all_signal_names
    assert "PCAP.BITS0.Value" not in all_signal_names
    assert "COUNTER1.OUT.Value" not in all_signal_names


def test_panda_csaxs(panda_csaxs):
    assert panda_csaxs.name == "panda_csaxs"
    assert panda_csaxs.host == "csaxs-panda-box.psi.ch"

    assert panda_csaxs._acquisition_group == "burst"
    assert panda_csaxs._timeout_on_completed == 10


@pytest.mark.parametrize(
    "scan_type, frames_per_trigger, expected_acquisition_group",
    [
        ("hardware_triggered", 1, "fly"),
        ("hardware_triggered", 5, "fly"),
        ("software_triggered", 10, "burst"),
        ("software_triggered", 1, "monitored"),  # Default case
    ],
)
def test_panda_csaxs_stage(panda_csaxs, scan_type, frames_per_trigger, expected_acquisition_group):
    """Test the on_stage method of the PandaBoxCSAXS device for different scan types and frames per trigger."""
    assert len(panda_csaxs._status_callbacks) == 0

    panda_csaxs.scan_info.msg.info["scan_type"] = scan_type
    panda_csaxs.scan_info.msg.info["frames_per_trigger"] = frames_per_trigger
    panda_csaxs.stage()

    assert panda_csaxs._acquisition_group == expected_acquisition_group
    assert panda_csaxs.staged == Staged.yes


def test_panda_csaxs_complete(panda_csaxs):
    """Test the on_complete method of the PandaBoxCSAXS device."""
    panda_csaxs.scan_info.msg.info["num_points"] = 1
    panda_csaxs.scan_info.msg.info["frames_per_trigger"] = 1
    panda_csaxs.scan_parameters = fetch_scan_info(panda_csaxs.scan_info)

    panda_csaxs._timeout_on_completed = 0.5  # Set a short timeout for testing

    def _mock_return_captured(*args, **kwargs):
        return ["=0"]

    # Timeout Error on complete
    with (
        mock.patch.object(panda_csaxs, "send_raw", side_effect=_mock_return_captured),
        mock.patch.object(panda_csaxs, "_disarm", return_value=None) as mock_disarm,
    ):
        status = panda_csaxs.on_complete()
        assert status.done is False
        assert status.success is False

        with pytest.raises(TimeoutError):
            status.wait(timeout=4)
        mock_disarm.assert_called_once()

    # Successful complete
    panda_csaxs._timeout_on_completed = 5
    with (
        mock.patch.object(panda_csaxs, "send_raw", side_effect=[["=0"], ["=0"], ["=1"]]),
        mock.patch.object(panda_csaxs, "_disarm", return_value=None) as mock_disarm,
    ):
        status = panda_csaxs.on_complete()
        assert status.done is False
        assert status.success is False

        status.wait(timeout=4)
        mock_disarm.assert_called_once()
        assert status.done is True
        assert status.success is True