185 lines
6.4 KiB
Python
185 lines
6.4 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
from aare.devices.jfjoch import JFJochWrapper
|
|
from aare.common.beamline import MXBeamline
|
|
from aare.common.rotation_scan import RotationScanRequest
|
|
from aare.common.raster_grid import RasterGridRequest
|
|
from aare.common.models import DAQStatusModel, SampleShortInfo, BeamlineStatus, SessionStatus, BeamlineStateEnum, SessionsStateEnum, SampleCameraSettings
|
|
from aare.common.diffraction_geometry import DiffractionGeometry
|
|
from aare.common.sample_geometry import SampleGeometryModel
|
|
from aare.common.coordinate import Coordinate, SmargonCoordinate
|
|
from aare.common.exception_handler import JFJochCommunicationError
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_jfjoch_client():
|
|
with patch('jfjoch_client.ApiClient'), \
|
|
patch('jfjoch_client.DefaultApi') as mock_api_class, \
|
|
patch('jfjoch_client.Configuration'):
|
|
mock_api = mock_api_class.return_value
|
|
yield mock_api
|
|
|
|
|
|
@pytest.fixture
|
|
def jfjoch_wrapper(mock_jfjoch_client):
|
|
return JFJochWrapper(MXBeamline.X10SA)
|
|
|
|
|
|
def create_mock_daq_status():
|
|
bl_status = BeamlineStatus(
|
|
name="test_beamline",
|
|
ring_current_mA=400.0,
|
|
front_light=50.0,
|
|
back_light=0.0,
|
|
cryojet_K=100.0,
|
|
shutter_open=False,
|
|
exp_shutter_open=False,
|
|
flux_ph_s=1e12,
|
|
sample_camera=SampleCameraSettings(exposure=0.001, gain=0.0),
|
|
transmission=1.0,
|
|
zoom=1.0,
|
|
commissioning_mode=False,
|
|
dtz_min=120.0,
|
|
dtz_max=1600.0
|
|
)
|
|
diff_geom = DiffractionGeometry(
|
|
energy_keV=12.658,
|
|
dtz_mm=200.0,
|
|
pixel_size_mm=0.075,
|
|
beam_center_pxl=(1000.0, 1000.0),
|
|
detector_size_pxl=(2000, 2000),
|
|
detector_description="Eiger",
|
|
detector_serial_number="123",
|
|
poni_rot1_rad=0.0,
|
|
poni_rot2_rad=0.0
|
|
)
|
|
session = SessionStatus(session=SessionsStateEnum.Vacant, current_pgroup="p12345")
|
|
|
|
geom = SampleGeometryModel(
|
|
beam_location_pxl=Coordinate(x=1000, y=1000),
|
|
pixel_in_mm=0.001,
|
|
aerotech=Coordinate(x=0, y=0),
|
|
aerotech_meas=Coordinate(x=0, y=0),
|
|
smargon=SmargonCoordinate(sh_mm=Coordinate(x=0, y=0), phi_deg=0, chi_deg=0),
|
|
omega_deg=0.0,
|
|
beam_size_mm=Coordinate(x=0.01, y=0.01)
|
|
)
|
|
|
|
return DAQStatusModel(
|
|
geom=geom,
|
|
diffraction=diff_geom,
|
|
bl=bl_status,
|
|
state=BeamlineStateEnum.SampleAlignment,
|
|
busy=False,
|
|
session=session
|
|
)
|
|
|
|
|
|
def test_init_x10sa(mock_jfjoch_client):
|
|
wrapper = JFJochWrapper(MXBeamline.X10SA)
|
|
assert wrapper._JFJochWrapper__url == "http://sls-gpu-002:8080"
|
|
|
|
|
|
def test_init_x06da(mock_jfjoch_client):
|
|
wrapper = JFJochWrapper(MXBeamline.X06DA)
|
|
assert wrapper._JFJochWrapper__url == "http://sls-gpu-001:8080"
|
|
|
|
|
|
def test_init_simulated(mock_jfjoch_client):
|
|
wrapper = JFJochWrapper(MXBeamline.SIMULATED)
|
|
assert wrapper._JFJochWrapper__simulated is True
|
|
|
|
|
|
def test_init_unknown():
|
|
with pytest.raises(Exception, match="unknown beamline"):
|
|
JFJochWrapper("UNKNOWN")
|
|
|
|
|
|
def test_initialize(jfjoch_wrapper, mock_jfjoch_client):
|
|
mock_jfjoch_client.initialize_post.reset_mock()
|
|
jfjoch_wrapper.initialize()
|
|
mock_jfjoch_client.initialize_post.assert_called_once()
|
|
|
|
|
|
def test_cancel(jfjoch_wrapper, mock_jfjoch_client):
|
|
mock_jfjoch_client.cancel_post.reset_mock()
|
|
jfjoch_wrapper.cancel()
|
|
mock_jfjoch_client.cancel_post.assert_called_once()
|
|
|
|
|
|
def test_is_idle(jfjoch_wrapper, mock_jfjoch_client):
|
|
mock_jfjoch_client.status_get.return_value.state = 'Idle'
|
|
assert jfjoch_wrapper.is_idle() is True
|
|
|
|
mock_jfjoch_client.status_get.return_value.state = 'Busy'
|
|
assert jfjoch_wrapper.is_idle() is False
|
|
|
|
|
|
def test_measure_rotation(jfjoch_wrapper, mock_jfjoch_client):
|
|
r = RotationScanRequest(
|
|
start_omega_deg=0.0,
|
|
incr_omega_deg=0.1,
|
|
steps=100,
|
|
exp_time_s=0.01,
|
|
file_prefix="test_prefix"
|
|
)
|
|
s = create_mock_daq_status()
|
|
s.sample = SampleShortInfo(sample_name="S1", user="p12345", db_id=1, dewar_name="D1", puck_name="P1", run_number=1, pin=1)
|
|
|
|
jfjoch_wrapper.measure_rotation(r, s)
|
|
mock_jfjoch_client.start_post.assert_called_once()
|
|
|
|
|
|
def test_measure_rotation_error(jfjoch_wrapper, mock_jfjoch_client):
|
|
mock_jfjoch_client.start_post.side_effect = Exception("API Error")
|
|
r = RotationScanRequest(exp_time_s=0.01, incr_omega_deg=0.1, steps=10)
|
|
s = create_mock_daq_status()
|
|
s.sample = SampleShortInfo(sample_name="S1", user="p12345", db_id=1, dewar_name="D1", puck_name="P1", run_number=1, pin=1)
|
|
|
|
with pytest.raises(JFJochCommunicationError):
|
|
jfjoch_wrapper.measure_rotation(r, s)
|
|
|
|
|
|
def test_measure_raster(jfjoch_wrapper, mock_jfjoch_client):
|
|
r = RasterGridRequest(
|
|
n_x=10,
|
|
n_y=10,
|
|
grid_size_mm=Coordinate(x=0.01, y=0.01),
|
|
exp_time_s=0.01,
|
|
file_prefix="raster_prefix",
|
|
smargon_top_left=SmargonCoordinate(sh_mm=Coordinate(x=0, y=0))
|
|
)
|
|
s = create_mock_daq_status()
|
|
s.sample = SampleShortInfo(sample_name="S1", user="p12345", db_id=1, dewar_name="D1", puck_name="P1", run_number=1, pin=1)
|
|
|
|
jfjoch_wrapper.measure_raster(r, s)
|
|
mock_jfjoch_client.start_post.assert_called_once()
|
|
|
|
|
|
def test_wait_till_done(jfjoch_wrapper, mock_jfjoch_client):
|
|
mock_jfjoch_client.result_scan_get.return_value = MagicMock()
|
|
result = jfjoch_wrapper.wait_till_done(timeout=10)
|
|
mock_jfjoch_client.wait_till_done_post_with_http_info.assert_called_once_with(timeout=10)
|
|
assert result is not None
|
|
|
|
|
|
def test_detector(jfjoch_wrapper, mock_jfjoch_client):
|
|
mock_det = MagicMock()
|
|
mock_jfjoch_client.config_select_detector_get.return_value.detectors = [mock_det]
|
|
mock_jfjoch_client.config_select_detector_get.return_value.current_id = 0
|
|
|
|
det = jfjoch_wrapper.detector()
|
|
assert det == mock_det
|
|
|
|
|
|
def test_get_diffraction_image(jfjoch_wrapper, mock_jfjoch_client):
|
|
mock_jfjoch_client.image_buffer_image_jpeg_get.return_value = b"image_data"
|
|
img = jfjoch_wrapper.get_diffraction_image(image_id=1)
|
|
assert img == b"image_data"
|
|
|
|
|
|
def test_get_diffraction_image_retry(jfjoch_wrapper, mock_jfjoch_client):
|
|
mock_jfjoch_client.image_buffer_image_jpeg_get.side_effect = [Exception("Error"), b"image_data"]
|
|
img = jfjoch_wrapper.get_diffraction_image(image_id=1, wait_between_retries_s=0.001)
|
|
assert img == b"image_data"
|
|
assert mock_jfjoch_client.image_buffer_image_jpeg_get.call_count == 2 |