WIP agent step 3: migrate scan tests to v4

2026-05-28 08:22:48 +02:00
parent 5bedd98292
commit 7d1e6d5bef
3 changed files with 204 additions and 0 deletions
@@ -0,0 +1,7 @@
+# pylint: skip-file
+from bec_server.scan_server.tests.scan_fixtures import (
+    nth_done_status_mock,
+    readout_priority,
+    v4_scan_assembler,
+)
+
@@ -0,0 +1,121 @@
+# pylint: skip-file
+from unittest import mock
+
+import numpy as np
+import pytest
+from bec_server.scan_server.tests.scan_fixtures import *
+from bec_server.scan_server.tests.scan_hook_tests import *
+
+XAS_SIMPLE_SCAN_DEFAULT_HOOK_TESTS = [
+    ("prepare_scan", [assert_prepare_scan_reads_baseline_devices]),
+    ("open_scan", [assert_scan_open_called]),
+    ("pre_scan", [assert_pre_scan_called]),
+    ("unstage", [assert_unstage_all_devices_called]),
+    ("close_scan", [assert_close_scan_waits_for_baseline_and_closes]),
+]
+
+
+def _assemble_xas_simple_scan(v4_scan_assembler, **overrides):
+    params = {
+        "start": 8000.0,
+        "stop": 9000.0,
+        "scan_time": 1.0,
+        "scan_duration": 10.0,
+        "motor": "mo1_bragg",
+        "daq": "nidaq",
+        "gonio": "mo1_gonio",
+        "monitored_readout_cycle": 1.0,
+    }
+    params.update(overrides)
+    return v4_scan_assembler("xas_simple_scan", **params)
+
+
+@pytest.mark.parametrize(("hook_name", "hook_tests"), XAS_SIMPLE_SCAN_DEFAULT_HOOK_TESTS)
+def test_xas_simple_scan_v4_default_hooks(
+    v4_scan_assembler, nth_done_status_mock, hook_name, hook_tests
+):
+    scan = _assemble_xas_simple_scan(v4_scan_assembler)
+
+    run_scan_tests(scan, [(hook_name, hook_tests)], nth_done_status_mock=nth_done_status_mock)
+
+
+def test_xas_simple_scan_v4_prepare_scan_updates_metadata(v4_scan_assembler):
+    scan = _assemble_xas_simple_scan(v4_scan_assembler)
+    scan.actions.add_scan_report_instruction_device_progress = mock.MagicMock()
+    baseline_status = mock.MagicMock()
+    scan.actions.read_baseline_devices = mock.MagicMock(return_value=baseline_status)
+
+    scan.prepare_scan()
+
+    scan.actions._build_scan_status_message("open")
+
+    np.testing.assert_array_equal(scan.scan_info.positions, np.array([8000.0, 9000.0]))
+    assert scan.scan_info.additional_scan_parameters["scan_time"] == 1.0
+    assert scan.scan_info.additional_scan_parameters["scan_duration"] == 10.0
+    assert scan.scan_info.additional_scan_parameters["monitored_readout_cycle"] == 1.0
+    assert scan.scan_info.readout_priority_modification["async"] == ["nidaq"]
+    scan.actions.add_scan_report_instruction_device_progress.assert_called_once_with(scan.motor)
+    scan.actions.read_baseline_devices.assert_called_once_with(wait=False)
+    assert scan._baseline_readout_status is baseline_status
+
+
+def test_xas_simple_scan_v4_stage_centers_gonio_and_stages_devices(v4_scan_assembler):
+    scan = _assemble_xas_simple_scan(v4_scan_assembler)
+    scan.mo1_bragg.convert_angle_energy = mock.MagicMock(side_effect=[10.0, 14.0])
+    set_status = mock.MagicMock()
+    scan.actions.set = mock.MagicMock(return_value=set_status)
+    scan.actions.stage_all_devices = mock.MagicMock()
+
+    scan.stage()
+
+    assert scan.mo1_bragg.convert_angle_energy.call_args_list == [
+        mock.call(mode="EnergyToAngle", inp=8000.0),
+        mock.call(mode="EnergyToAngle", inp=9000.0),
+    ]
+    scan.actions.set.assert_called_once_with(scan.mo1_gonio, 12.0, wait=False)
+    set_status.wait.assert_called_once_with(timeout=30)
+    scan.actions.stage_all_devices.assert_called_once_with()
+
+
+def test_xas_simple_scan_v4_scan_core_reads_until_complete(v4_scan_assembler, nth_done_status_mock):
+    scan = _assemble_xas_simple_scan(v4_scan_assembler)
+    completion_status = nth_done_status_mock(resolve_after=3)
+    scan.actions.kickoff = mock.MagicMock()
+    scan.actions.complete = mock.MagicMock(return_value=completion_status)
+    scan.actions.read_monitored_devices = mock.MagicMock()
+
+    with mock.patch("superxas_bec.scans.mono_bragg_scans.time.sleep"):
+        scan.scan_core()
+
+    scan.actions.kickoff.assert_called_once_with(scan.motor)
+    scan.actions.complete.assert_called_once_with(scan.motor, wait=False)
+    assert scan.actions.read_monitored_devices.call_count == 2
+
+
+def test_xas_simple_scan_v4_post_scan_completes_all_devices(v4_scan_assembler):
+    scan = _assemble_xas_simple_scan(v4_scan_assembler)
+    scan.actions.complete_all_devices = mock.MagicMock()
+
+    scan.post_scan()
+
+    scan.actions.complete_all_devices.assert_called_once_with()
+
+
+def test_xas_advanced_scan_v4_updates_spline_metadata(v4_scan_assembler):
+    scan = v4_scan_assembler(
+        "xas_advanced_scan",
+        start=8000.0,
+        stop=9000.0,
+        scan_time=1.0,
+        scan_duration=10.0,
+        p_kink=50.0,
+        e_kink=8500.0,
+        motor="mo1_bragg",
+        daq="nidaq",
+        gonio="mo1_gonio",
+    )
+
+    assert scan.scan_name == "xas_advanced_scan"
+    assert scan.scan_info.additional_scan_parameters["p_kink"] == 50.0
+    assert scan.scan_info.additional_scan_parameters["e_kink"] == 8500.0
+
@@ -0,0 +1,76 @@
+# pylint: skip-file
+from unittest import mock
+
+import pytest
+from bec_server.scan_server.tests.scan_fixtures import *
+from bec_server.scan_server.tests.scan_hook_tests import *
+
+NIDAQ_CONTINUOUS_SCAN_DEFAULT_HOOK_TESTS = [
+    ("prepare_scan", [assert_prepare_scan_reads_baseline_devices]),
+    ("open_scan", [assert_scan_open_called]),
+    ("stage", [assert_stage_all_devices_called]),
+    ("pre_scan", [assert_pre_scan_called]),
+    ("unstage", [assert_unstage_all_devices_called]),
+    ("close_scan", [assert_close_scan_waits_for_baseline_and_closes]),
+]
+
+
+def _assemble_nidaq_continuous_scan(v4_scan_assembler, **overrides):
+    params = {"scan_duration": 10.0, "daq": "nidaq", "compression": False}
+    params.update(overrides)
+    return v4_scan_assembler("nidaq_continuous_scan", **params)
+
+
+@pytest.mark.parametrize(("hook_name", "hook_tests"), NIDAQ_CONTINUOUS_SCAN_DEFAULT_HOOK_TESTS)
+def test_nidaq_continuous_scan_v4_default_hooks(
+    v4_scan_assembler, nth_done_status_mock, hook_name, hook_tests
+):
+    scan = _assemble_nidaq_continuous_scan(v4_scan_assembler)
+
+    run_scan_tests(scan, [(hook_name, hook_tests)], nth_done_status_mock=nth_done_status_mock)
+
+
+def test_nidaq_continuous_scan_v4_prepare_scan_updates_metadata(v4_scan_assembler):
+    scan = _assemble_nidaq_continuous_scan(v4_scan_assembler)
+    scan.actions.add_scan_report_instruction_device_progress = mock.MagicMock()
+    baseline_status = mock.MagicMock()
+    scan.actions.read_baseline_devices = mock.MagicMock(return_value=baseline_status)
+
+    scan.prepare_scan()
+    scan.actions._build_scan_status_message("open")
+
+    assert scan.scan_info.additional_scan_parameters["scan_duration"] == 10.0
+    assert scan.scan_info.additional_scan_parameters["compression"] is False
+    assert scan.scan_info.readout_priority_modification["async"] == ["nidaq"]
+    scan.actions.add_scan_report_instruction_device_progress.assert_called_once_with(scan.daq)
+    scan.actions.read_baseline_devices.assert_called_once_with(wait=False)
+    assert scan._baseline_readout_status is baseline_status
+
+
+def test_nidaq_continuous_scan_v4_scan_core_reads_until_complete(
+    v4_scan_assembler, nth_done_status_mock
+):
+    scan = _assemble_nidaq_continuous_scan(v4_scan_assembler)
+    kickoff_status = mock.MagicMock()
+    completion_status = nth_done_status_mock(resolve_after=3)
+    scan.actions.kickoff = mock.MagicMock(return_value=kickoff_status)
+    scan.actions.complete = mock.MagicMock(return_value=completion_status)
+    scan.actions.read_monitored_devices = mock.MagicMock()
+
+    with mock.patch("superxas_bec.scans.nidaq_cont_scan.time.sleep"):
+        scan.scan_core()
+
+    scan.actions.kickoff.assert_called_once_with(device=scan.daq, wait=False)
+    kickoff_status.wait.assert_called_once_with(timeout=5)
+    scan.actions.complete.assert_called_once_with(device=scan.daq, wait=False)
+    assert scan.actions.read_monitored_devices.call_count == 2
+
+
+def test_nidaq_continuous_scan_v4_post_scan_completes_all_devices(v4_scan_assembler):
+    scan = _assemble_nidaq_continuous_scan(v4_scan_assembler)
+    scan.actions.complete_all_devices = mock.MagicMock()
+
+    scan.post_scan()
+
+    scan.actions.complete_all_devices.assert_called_once_with()
+