csaxs_bec/tests/tests_devices/test_falcon_csaxs.py

# pylint: skip-file
import inspect
import os
import threading
from typing import Generator
from unittest import mock

import ophyd
import pytest
from bec_lib import messages
from bec_lib.endpoints import MessageEndpoints
from bec_lib.file_utils import get_full_path
from bec_server.device_server.tests.utils import DMMock
from ophyd_devices.interfaces.base_classes.psi_device_base import DeviceStoppedError
from ophyd_devices.tests.utils import patched_device

from csaxs_bec.devices.epics.falcon_csaxs import (
    ACQUIRESTATUS,
    DEFAULT_XML_FILE_NAME,
    FalconcSAXS,
    FalconError,
    MappingSource,
    TriggerSource,
)
from csaxs_bec.devices.utils.utils import fetch_scan_info


@pytest.fixture(scope="function")
def mock_det() -> Generator[FalconcSAXS, None, None]:
    """Fixture to mock the FalconcSAXS device."""
    name = "mcs_csaxs"
    prefix = "X12SA-MCS-CSAXS:"
    dm = DMMock()
    with patched_device(
        FalconcSAXS,
        name="falcon",
        prefix="X12SA-SITORO:",
        device_manager=dm,
        _mock_pv_initial_value=1,
    ) as dev:
        try:
            for dotted_name, device in dev.walk_subdevices(include_lazy=True):
                device.stage_sigs = {}  # Remove stage signals
                device.trigger_sigs = {}  # Remove trigger signals
                if hasattr(device, "plugin_type"):
                    device.plugin_type._read_pv.mock_data = device._plugin_type
            yield dev
        finally:
            dev.destroy()


def test_falcon_init(mock_det: FalconcSAXS):
    """Test the initialization of the FalconcSAXS device."""
    assert mock_det._readout_time == mock_det.MIN_READOUT
    assert mock_det._value_pixel_per_buffer == 20
    assert mock_det._queue_size == 2000
    assert mock_det._full_path == ""


def test_falcon_on_connected(mock_det: FalconcSAXS):
    """Test the on_connected method of the FalconcSAXS device."""
    falcon = mock_det

    # Set known default values
    falcon.preset_mode.put(-1)
    falcon.input_logic_polarity.put(-1)
    falcon.auto_pixels_per_buffer.put(-1)
    falcon.hdf5.enable.put(-1)

    with (
        mock.patch.object(falcon, "on_stop") as mock_on_stop,
        mock.patch.object(falcon, "set_trigger") as mock_set_trigger,
    ):

        falcon.on_connected()
        mock_on_stop.assert_called_once()
        mock_set_trigger.assert_called_once_with(
            mapping_mode=MappingSource.MAPPING, trigger_source=TriggerSource.GATE, ignore_gate=0
        )

        # Detector default PV values
        assert falcon.preset_mode.get() == "1"  # Real Time
        assert falcon.input_logic_polarity.get() == 0
        assert falcon.auto_pixels_per_buffer.get() == 0
        assert falcon.pixels_per_buffer.get() == falcon._value_pixel_per_buffer

        # Backend default PV values
        assert falcon.hdf5.enable.get() == "1"  # Enabled
        assert falcon.hdf5.xml_file_name.get() == DEFAULT_XML_FILE_NAME
        assert falcon.hdf5.lazy_open.get() == "1"  # Enabled
        assert falcon.hdf5.temp_suffix.get() == ""
        assert falcon.hdf5.queue_size.get() == falcon._queue_size
        assert falcon.nd_array_mode.get() == 1
        assert falcon.hdf5.file_template.get() == "%s%s"
        assert falcon.hdf5.file_write_mode.get() == 2


def test_falcon_on_stage(mock_det: FalconcSAXS):
    """

    Test the on_stage method of the FalconcSAXS device.
    All relevant information is available in the scan_info attribute and used
    to bootstrap the detector for the upcoming acquisition. Two scenarios are tested:
    I. Normal case with exposure time larger than readout time
    II. Case where exposure time is smaller than readout time, which should raise an exception.
    """
    falcon = mock_det
    num_points = 10
    exp_time = 0.2
    frames_per_trigger = 5
    falcon.scan_info.msg.info["num_points"] = num_points
    falcon.scan_info.msg.info["frames_per_trigger"] = frames_per_trigger
    falcon.scan_info.msg.info["exp_time"] = exp_time
    falcon.hdf5.array_counter.put(5)  # Set to non-zero to check reset

    # I. Normal case
    falcon.stage()

    assert falcon.staged is ophyd.Staged.yes
    assert falcon._full_path == get_full_path(falcon.scan_info.msg, falcon.name)
    file_path = falcon.hdf5.file_path.get()
    file_name = falcon.hdf5.file_name.get()
    assert os.path.join(file_path, file_name) == falcon._full_path

    assert falcon.preset_real_time.get() == exp_time
    assert falcon.pixels_per_run.get() == num_points * frames_per_trigger
    assert falcon.hdf5.num_capture.get() == num_points * frames_per_trigger
    assert falcon.hdf5.array_counter.get() == 0
    assert falcon.hdf5.capture.get() == 1
    assert falcon.start_all.get() == 1

    # II. Unstage device first
    falcon.unstage()
    exp_time = 1e-3  # Smaller than readout time
    falcon.scan_info.msg.info["exp_time"] = exp_time
    with pytest.raises(ValueError):
        falcon.stage()
        assert falcon.staged is not ophyd.Staged.no


def test_falcon_on_pre_scan(mock_det: FalconcSAXS):
    """Test the on_pre_scan method of the FalconcSAXS device."""
    falcon = mock_det
    # I. Test normal case with success
    falcon.acquire_busy._read_pv.mock_data = ACQUIRESTATUS.DONE
    falcon.hdf5.capture._read_pv.mock_data = ACQUIRESTATUS.DONE
    falcon = mock_det
    st = falcon.on_pre_scan()
    assert st.done is False
    assert st.success is False
    falcon.acquire_busy._read_pv.mock_data = ACQUIRESTATUS.ACQUIRING
    assert st.done is False
    assert st.success is False
    falcon.hdf5.capture._read_pv.mock_data = ACQUIRESTATUS.ACQUIRING
    st.wait(3)
    assert st.done is True
    assert st.success is True

    # II. Test abort case with stop called
    falcon.acquire_busy._read_pv.mock_data = ACQUIRESTATUS.DONE
    falcon.hdf5.capture._read_pv.mock_data = ACQUIRESTATUS.DONE
    st = falcon.on_pre_scan()
    assert st.done is False
    assert st.success is False
    falcon.stop()
    with pytest.raises(DeviceStoppedError):
        st.wait(3)
    assert st.done is True
    assert st.success is False


def test_falcon_stop(mock_det: FalconcSAXS):
    """Test the stop method of the FalconcSAXS device."""
    falcon = mock_det

    falcon.stop_all.put(0)
    falcon.hdf5.capture.put(1)
    falcon.erase_all.put(0)
    falcon.stop()
    assert falcon.stop_all.get() == 1
    assert falcon.hdf5.capture.get() == 0
    assert falcon.erase_all.get() == 1


def test_falcon_complete(mock_det: FalconcSAXS):
    """Test the complete method of the FalconcSAXS device."""
    falcon = mock_det

    num_points = 10
    frames_per_trigger = 5
    falcon.scan_info.msg.info["num_points"] = num_points
    falcon.scan_info.msg.info["frames_per_trigger"] = frames_per_trigger
    falcon.scan_parameters = fetch_scan_info(falcon.scan_info)

    # I. Test normal case with success
    falcon.dxp.current_pixel._read_pv.mock_data = num_points * frames_per_trigger - 1
    falcon.hdf5.array_counter._read_pv.mock_data = num_points * frames_per_trigger - 1
    falcon._full_path = "/tmp/fake_path/test.h5"
    st = falcon.on_complete()
    assert st.done is False
    assert st.success is False
    falcon.dxp.current_pixel._read_pv.mock_data = num_points * frames_per_trigger
    assert st.done is False
    assert st.success is False
    falcon.hdf5.array_counter._read_pv.mock_data = num_points * frames_per_trigger
    st.wait(3)
    assert st.done is True
    assert st.success is True
    assert falcon.file_event.get() == messages.FileMessage(
        file_path="/tmp/fake_path/test.h5",
        done=True,
        successful=True,
        device_name=falcon.name,
        file_type="h5",
        hinted_h5_entries=None,
        metadata={},
    )

    # II. Test case where acquisition fails due to interruption
    falcon.dxp.current_pixel._read_pv.mock_data = num_points * frames_per_trigger - 1
    st = falcon.on_complete()
    assert st.done is False
    assert st.success is False
    falcon.stop()
    with pytest.raises(DeviceStoppedError):
        st.wait(3)
    assert falcon.file_event.get() == messages.FileMessage(
        file_path="/tmp/fake_path/test.h5",
        done=True,
        successful=False,
        device_name=falcon.name,
        file_type="h5",
        hinted_h5_entries=None,
        metadata={},
    )


def test_falcon_is_primed(mock_det: FalconcSAXS):
    """The plugin counts as primed as soon as any array dimension is non-zero."""
    falcon = mock_det

    with mock.patch.object(falcon.hdf5.array_size, "get", return_value=(0, 0, 0)):
        assert falcon.is_primed() is False

    with mock.patch.object(falcon.hdf5.array_size, "get", return_value=(0, 1, 3000)):
        assert falcon.is_primed() is True


def test_falcon_prime_skips_when_already_primed(mock_det: FalconcSAXS):
    """Priming an already primed plugin must not touch the detector."""
    falcon = mock_det

    with (
        mock.patch.object(falcon, "is_primed", return_value=True),
        mock.patch.object(falcon, "set_trigger") as mock_set_trigger,
        mock.patch.object(falcon, "on_stop") as mock_on_stop,
    ):
        falcon.prime()
        mock_set_trigger.assert_not_called()
        mock_on_stop.assert_not_called()


def test_falcon_prime(mock_det: FalconcSAXS):
    """Priming advances a pixel by hand and restores the trigger configuration."""
    falcon = mock_det

    falcon.set_trigger(
        mapping_mode=MappingSource.MAPPING, trigger_source=TriggerSource.GATE, ignore_gate=0
    )
    falcon.pixels_per_buffer.put(20)

    with (
        mock.patch.object(falcon, "is_primed", side_effect=[False, True]),
        mock.patch.object(falcon.next_pixel, "put") as mock_next_pixel,
    ):
        falcon.prime()

        # A gate signal must not be required, so the pixel is advanced from software.
        assert mock_next_pixel.call_count > 0
        assert falcon.hdf5.capture.get() == 0

    # Trigger configuration restored, otherwise the next scan silently starves.
    assert falcon.collect_mode.get() == MappingSource.MAPPING
    assert falcon.pixel_advance_mode.get() == TriggerSource.GATE
    assert falcon.ignore_gate.get() == 0
    assert falcon.pixels_per_buffer.get() == 20


def test_falcon_prime_restores_config_on_failure(mock_det: FalconcSAXS):
    """A failed priming must still leave the detector in its original trigger config."""
    falcon = mock_det

    falcon.set_trigger(
        mapping_mode=MappingSource.MAPPING, trigger_source=TriggerSource.GATE, ignore_gate=0
    )
    falcon.pixels_per_buffer.put(20)

    with (
        mock.patch.object(falcon, "is_primed", return_value=False),
        mock.patch.object(falcon, "_wait_for_primed", return_value=False),
    ):
        with pytest.raises(FalconError):
            falcon.prime()

    assert falcon.pixel_advance_mode.get() == TriggerSource.GATE
    assert falcon.ignore_gate.get() == 0
    assert falcon.pixels_per_buffer.get() == 20


def test_falcon_sets_the_ioc_resolvable_layout(mock_det: FalconcSAXS):
    """A bare 'layout.xml' does not resolve for the IOC; cfg/layout.xml does."""
    falcon = mock_det
    falcon.hdf5.xml_file_name.put("something-else")

    with mock.patch.object(falcon, "_check_xml_layout", return_value=True):
        falcon._initialize_detector_backend()

    assert falcon.hdf5.xml_file_name.get() == DEFAULT_XML_FILE_NAME
    assert DEFAULT_XML_FILE_NAME == "cfg/layout.xml"


def test_falcon_xml_file_name_from_device_config():
    """deviceConfig can point at a different layout, or disable it with ''."""
    with patched_device(
        FalconcSAXS,
        name="falcon",
        prefix="X12SA-SITORO:",
        device_manager=DMMock(),
        xml_file_name="",
        _mock_pv_initial_value=1,
    ) as dev:
        try:
            assert dev._xml_file_name == ""
            # Nothing to validate when the built-in layout is requested.
            assert dev._check_xml_layout(timeout=0.01) is True
        finally:
            dev.destroy()


def test_falcon_warns_when_the_ioc_rejects_the_layout(mock_det: FalconcSAXS):
    """An unreadable layout is reported at connect, not as a write error much later."""
    falcon = mock_det

    with (
        mock.patch.object(falcon.hdf5.xml_valid, "get", return_value="No"),
        mock.patch.object(falcon.hdf5.xml_error_msg, "get", return_value="cannot be opened"),
        mock.patch("csaxs_bec.devices.epics.falcon_csaxs.logger") as mock_logger,
    ):
        assert falcon._check_xml_layout(timeout=0.05) is False
        assert mock_logger.error.call_count == 1
        assert "layout" in mock_logger.error.call_args[0][0].lower()


def test_falcon_accepts_a_valid_layout_silently(mock_det: FalconcSAXS):
    """A layout the IOC accepts produces no error."""
    falcon = mock_det

    with (
        mock.patch.object(falcon.hdf5.xml_valid, "get", return_value="Yes"),
        mock.patch("csaxs_bec.devices.epics.falcon_csaxs.logger") as mock_logger,
    ):
        assert falcon._check_xml_layout(timeout=0.05) is True
        mock_logger.error.assert_not_called()


def test_falcon_init_names_every_device_config_key():
    """The device server drops config keys that are not named parameters.

    bec_server's devicemanager builds the init kwargs by intersecting the deviceConfig
    keys with the class signature's named parameters. A *args/**kwargs signature
    therefore silently discards 'prefix', and every PV is built without it -- which
    presents as 'failed to connect to all signals' rather than as a config error.
    """
    params = inspect.signature(FalconcSAXS).parameters
    for key in ("name", "prefix", "scan_info", "device_manager", "xml_file_name"):
        assert key in params, f"{key} must be a named parameter or the device server drops it"


def test_falcon_keeps_the_configured_prefix():
    """A prefix passed as the device server passes it must reach the signals."""
    with patched_device(
        FalconcSAXS,
        name="falcon",
        prefix="X12SA-SITORO:",
        device_manager=DMMock(),
        _mock_pv_initial_value=1,
    ) as dev:
        try:
            assert dev.prefix == "X12SA-SITORO:"
            assert dev.hdf5.file_path.pvname.startswith("X12SA-SITORO:")
        finally:
            dev.destroy()