from unittest.mock import MagicMock, patch import pytest from aarecommon.errors.exception_handler import JFJochCommunicationError from aarecommon.math.coordinate import Coordinate, SmargonCoordinate from aarecommon.math.diffraction_geometry import DiffractionGeometry from aarecommon.math.sample_geometry import SampleGeometryModel from aarecommon.models.beamline import MXBeamline from aarecommon.models.models import ( BeamlineStateEnum, BeamlineStatus, DAQStatusModel, SampleCameraSettings, SampleShortInfo, SessionsStateEnum, SessionStatus, ) from aarecommon.models.raster_grid import RasterGridRequest from aarecommon.models.rotation_scan import RotationScanRequest from aare.devices.jfjoch import JFJochWrapper @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._url == "http://sls-gpu-002:8080" def test_init_x06da(mock_jfjoch_client): wrapper = JFJochWrapper(MXBeamline.X06DA) assert wrapper._url == "http://sls-gpu-001:8080" def test_init_simulated(mock_jfjoch_client): wrapper = JFJochWrapper(MXBeamline.SIMULATED) assert wrapper._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", dtz=200.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_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, dtz=200.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, ) 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)), dtz=200.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