289 lines
12 KiB
Python
289 lines
12 KiB
Python
import pytest
|
|
import numpy as np
|
|
from unittest.mock import MagicMock, patch
|
|
from aare.daq.aaredb import AareWrapper
|
|
from aare.common.models import (
|
|
SampleShortInfo, PuckLoadedInfo, DewarAddress
|
|
)
|
|
from aare.common.raster_grid import RasterGridRequest, CenterOfMassModel
|
|
from aare.common.rotation_scan import RotationScanRequest
|
|
from aare.common.sample_geometry import SampleGeometryModel
|
|
from aare.common.coordinate import Coordinate, SmargonCoordinate
|
|
from jfjoch_client.models import ScanResult
|
|
|
|
@pytest.fixture
|
|
def mock_bl():
|
|
bl = MagicMock()
|
|
bl.value = "X10SA"
|
|
bl.name = "X10SA"
|
|
return bl
|
|
|
|
@pytest.fixture
|
|
def sample_info():
|
|
return SampleShortInfo(
|
|
db_id=123,
|
|
puck_name="P1",
|
|
dewar_name="D1",
|
|
sample_name="S1",
|
|
run_number=1,
|
|
user="testuser",
|
|
pin=5,
|
|
location=DewarAddress(segment="A", pos=1)
|
|
)
|
|
|
|
@pytest.fixture
|
|
def daq_status(mock_bl):
|
|
status = MagicMock()
|
|
status.bl = mock_bl
|
|
status.diffraction.detector_description = "EIGER2 16M"
|
|
status.diffraction.detector_serial_number = "1234"
|
|
status.diffraction.dtz_mm = 200.0
|
|
status.diffraction.max_resolution_angstrom = 1.5
|
|
status.diffraction.beam_center_pxl = [1000.0, 1000.0]
|
|
status.diffraction.pixel_size_mm = 0.075
|
|
status.diffraction.energy_keV = 12.658
|
|
status.diffraction.wavelength_angstrom = 0.979
|
|
status.bl.ring_current_mA = 400.0
|
|
status.bl.transmission = 1.0
|
|
status.bl.flux_ph_s = 1e12
|
|
status.bl.cryojet_K = 100.0
|
|
status.geom.beam_size_mm = Coordinate(x=0.01, y=0.01)
|
|
status.geom.smargon = SmargonCoordinate(sh_mm=Coordinate(x=0,y=0,z=0), phi_deg=0, chi_deg=0)
|
|
return status
|
|
|
|
@pytest.fixture
|
|
def geom_model():
|
|
return 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,z=0), phi_deg=0, chi_deg=0),
|
|
omega_deg=0,
|
|
beam_size_mm=Coordinate(x=0.01, y=0.01)
|
|
)
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("aareDB.TellsRunnerApi")
|
|
@patch("aareDB.SamplesRunnerApi")
|
|
@patch("aareDB.ProcessingsRunnerApi")
|
|
@patch("aareDB.GridscanRunnerApi")
|
|
def test_aare_wrapper_init(mock_grid, mock_proc, mock_sample, mock_tell, mock_api, mock_bl):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
mock_api.assert_called_once()
|
|
assert wrapper._AareWrapper__bl == mock_bl
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("aareDB.TellsRunnerApi")
|
|
@patch("aareDB.SamplesRunnerApi")
|
|
@patch("aareDB.ProcessingsRunnerApi")
|
|
@patch("aareDB.GridscanRunnerApi")
|
|
def test_set_pucks_beamline(mock_grid, mock_proc, mock_sample, mock_tell, mock_api, mock_bl):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
pucks = [PuckLoadedInfo(puck_name="P1", location=DewarAddress(segment="A", pos=1))]
|
|
wrapper.set_pucks_beamline(pucks)
|
|
mock_tell.return_value.set_tell_positions.assert_called_once()
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("aareDB.SamplesRunnerApi")
|
|
def test_create_manual_sample(mock_sample, mock_api, mock_bl, sample_info):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
mock_sample.return_value.insert_sample.return_value.id = 456
|
|
wrapper.create_manual_sample(sample_info)
|
|
assert sample_info.db_id == 456
|
|
mock_sample.return_value.insert_sample.assert_called_once()
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("aareDB.SamplesRunnerApi")
|
|
def test_send_sample_event(mock_sample, mock_api, mock_bl, sample_info):
|
|
from aareDB import SampleEventType
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
wrapper.send_sample_event(sample_info, SampleEventType.MOUNTED, "Test comment")
|
|
mock_sample.return_value.create_sample_event.assert_called_once()
|
|
|
|
# Test None sample
|
|
mock_sample.return_value.create_sample_event.reset_mock()
|
|
wrapper.send_sample_event(None, SampleEventType.MOUNTED)
|
|
mock_sample.return_value.create_sample_event.assert_not_called()
|
|
|
|
# Test invalid ID
|
|
sample_info.db_id = -1
|
|
wrapper.send_sample_event(sample_info, SampleEventType.MOUNTED)
|
|
mock_sample.return_value.create_sample_event.assert_not_called()
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("aareDB.SamplesRunnerApi")
|
|
def test_create_manual_sample_error(mock_sample, mock_api, mock_bl, sample_info, caplog):
|
|
import logging
|
|
caplog.set_level(logging.ERROR)
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
mock_sample.return_value.insert_sample.side_effect = Exception("DB Error")
|
|
wrapper.create_manual_sample(sample_info)
|
|
assert "Error inserting sample: DB Error" in caplog.text
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("aareDB.SamplesRunnerApi")
|
|
def test_send_sample_event_error(mock_sample, mock_api, mock_bl, sample_info, caplog):
|
|
from aareDB import SampleEventType
|
|
import logging
|
|
caplog.set_level(logging.ERROR)
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
mock_sample.return_value.create_sample_event.side_effect = Exception("Event Error")
|
|
wrapper.send_sample_event(sample_info, SampleEventType.MOUNTED)
|
|
assert "Error sending sample event" in caplog.text
|
|
assert "MOUNTED" in caplog.text
|
|
assert "Event Error" in caplog.text
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("requests.post")
|
|
def test_upload_image(mock_post, mock_api, mock_bl):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
mock_post.return_value.status_code = 200
|
|
img = np.zeros((100, 100, 3), dtype=np.uint8)
|
|
wrapper.upload_image(123, "test_img", img, "test message")
|
|
mock_post.assert_called_once()
|
|
assert "test_img.jpg" in str(mock_post.call_args)
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("requests.post")
|
|
def test_upload_jpg(mock_post, mock_api, mock_bl):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
mock_post.return_value.status_code = 200
|
|
wrapper.upload_jpg(123, "test_img", b"fake jpeg bytes")
|
|
mock_post.assert_called_once()
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("aareDB.SamplesRunnerApi")
|
|
def test_create_rotation_run(mock_sample, mock_api, mock_bl, sample_info, daq_status):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
# Standard rotation
|
|
req = RotationScanRequest(exp_time_s=0.1, incr_omega_deg=0.1, steps=100, file_prefix="test_prefix")
|
|
wrapper.create_rotation_run(sample_info, req, daq_status)
|
|
mock_sample.return_value.create_experiment_parameters_for_sample.assert_called_once()
|
|
|
|
# None sample - should return early
|
|
mock_sample.return_value.create_experiment_parameters_for_sample.reset_mock()
|
|
wrapper.create_rotation_run(None, req, daq_status)
|
|
mock_sample.return_value.create_experiment_parameters_for_sample.assert_not_called()
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("aareDB.SamplesRunnerApi")
|
|
def test_create_rotation_run_screening(mock_sample, mock_api, mock_bl, sample_info, daq_status):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
req = RotationScanRequest(exp_time_s=0.1, incr_omega_deg=0.1, steps=100, file_prefix="test_prefix", screening=True, wedge_omega_deg=5.0)
|
|
wrapper.create_rotation_run(sample_info, req, daq_status)
|
|
mock_sample.return_value.create_experiment_parameters_for_sample.assert_called_once()
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("aareDB.SamplesRunnerApi")
|
|
def test_create_rotation_run_error(mock_sample, mock_api, mock_bl, sample_info, daq_status, caplog):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
req = RotationScanRequest(exp_time_s=0.1, incr_omega_deg=0.1, steps=100, file_prefix="test_prefix")
|
|
mock_sample.return_value.create_experiment_parameters_for_sample.side_effect = Exception("Rotation error")
|
|
wrapper.create_rotation_run(sample_info, req, daq_status)
|
|
assert "Rotation error" in caplog.text
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("aareDB.SamplesRunnerApi")
|
|
def test_create_gridscan_run(mock_sample, mock_api, mock_bl, sample_info, daq_status):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
req = RasterGridRequest(exp_time_s=0.1, n_x=10, n_y=10, grid_size_mm=Coordinate(x=0.01, y=0.01), smargon_top_left=None, file_prefix="test_grid_prefix")
|
|
wrapper.create_gridscan_run(sample_info, req, daq_status)
|
|
mock_sample.return_value.create_experiment_parameters_for_sample.assert_called_once()
|
|
|
|
# None sample
|
|
mock_sample.return_value.create_experiment_parameters_for_sample.reset_mock()
|
|
wrapper.create_gridscan_run(None, req, daq_status)
|
|
mock_sample.return_value.create_experiment_parameters_for_sample.assert_not_called()
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("aareDB.SamplesRunnerApi")
|
|
def test_create_gridscan_run_error(mock_sample, mock_api, mock_bl, sample_info, daq_status, caplog):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
req = RasterGridRequest(exp_time_s=0.1, n_x=10, n_y=10, grid_size_mm=Coordinate(x=0.01, y=0.01), smargon_top_left=None, file_prefix="test_grid_prefix")
|
|
mock_sample.return_value.create_experiment_parameters_for_sample.side_effect = Exception("Grid error")
|
|
wrapper.create_gridscan_run(sample_info, req, daq_status)
|
|
assert "Grid error" in caplog.text
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("requests.post")
|
|
def test_ingest_gridscan(mock_post, mock_api, mock_bl, sample_info, geom_model):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
mock_post.return_value.status_code = 200
|
|
raster_result = MagicMock(spec=ScanResult)
|
|
raster_request = RasterGridRequest(exp_time_s=0.1, n_x=10, n_y=10, grid_size_mm=Coordinate(x=0.01, y=0.01), smargon_top_left=None)
|
|
com = CenterOfMassModel(n_x=5.0, n_y=5.0)
|
|
|
|
wrapper.ingest_gridscan(sample_info, raster_result, raster_request, geom_model, com, (500.0, 500.0))
|
|
mock_post.assert_called_once()
|
|
|
|
# None sample
|
|
mock_post.reset_mock()
|
|
wrapper.ingest_gridscan(None, raster_result, raster_request, geom_model, com, (500.0, 500.0))
|
|
mock_post.assert_not_called()
|
|
|
|
@patch("aareDB.ApiClient")
|
|
@patch("requests.post")
|
|
def test_ingest_scan(mock_post, mock_api, mock_bl, sample_info, geom_model):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
mock_post.return_value.status_code = 200
|
|
result = MagicMock(spec=ScanResult)
|
|
|
|
wrapper.ingest_scan(sample_info, result, geom_model, (500.0, 500.0))
|
|
mock_post.assert_called_once()
|
|
|
|
# None sample
|
|
mock_post.reset_mock()
|
|
wrapper.ingest_scan(None, result, geom_model, (500.0, 500.0))
|
|
mock_post.assert_not_called()
|
|
|
|
@patch("aareDB.ApiClient")
|
|
def test_format_gridscan_payload_no_com(mock_api, mock_bl, sample_info, geom_model):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
raster_result = MagicMock(spec=ScanResult)
|
|
raster_request = RasterGridRequest(exp_time_s=0.1, n_x=10, n_y=10, grid_size_mm=Coordinate(x=0.01, y=0.01), smargon_top_left=None)
|
|
|
|
payload = wrapper.format_gridscan_payload(sample_info, raster_result, raster_request, geom_model, None, (500.0, 500.0))
|
|
assert payload.center_pxl is None
|
|
assert payload.sample_id == sample_info.db_id
|
|
|
|
@patch("aareDB.ApiClient")
|
|
def test_format_gridscan_payload_with_top_left(mock_api, mock_bl, sample_info, geom_model):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
raster_result = MagicMock(spec=ScanResult)
|
|
top_left = SmargonCoordinate(sh_mm=Coordinate(x=0.1, y=0.1, z=0.1), phi_deg=0, chi_deg=0)
|
|
raster_request = RasterGridRequest(exp_time_s=0.1, n_x=10, n_y=10, grid_size_mm=Coordinate(x=0.01, y=0.01), smargon_top_left=top_left)
|
|
|
|
payload = wrapper.format_gridscan_payload(sample_info, raster_result, raster_request, geom_model, None, (500.0, 500.0))
|
|
assert payload.sample_id == sample_info.db_id
|
|
|
|
@patch("aareDB.ApiClient")
|
|
def test_ingest_gridscan_payload_none(mock_api, mock_bl, sample_info, geom_model):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
# We want format_gridscan_payload to return None
|
|
# Looking at aaredb.py, it returns None if an exception occurs and it's caught
|
|
# But wait, it raises 'e' after logging. Oh, I see.
|
|
# Lines 351-352 in aaredb.py: if payload is None: return
|
|
# This happens if format_gridscan_payload returns None.
|
|
with patch.object(AareWrapper, 'format_gridscan_payload', return_value=None):
|
|
wrapper.ingest_gridscan(sample_info, None, None, geom_model, None, (0,0))
|
|
|
|
with patch.object(AareWrapper, 'format_scan_payload', return_value=None):
|
|
wrapper.ingest_scan(sample_info, None, geom_model, (0,0))
|
|
|
|
@patch("aareDB.ApiClient")
|
|
def test_format_gridscan_payload_error(mock_api, mock_bl, sample_info, caplog):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
# Passing None for geom_model should trigger an error in smargon_to_picture
|
|
with pytest.raises(Exception):
|
|
wrapper.format_gridscan_payload(sample_info, None, None, None, None, (0,0))
|
|
assert "NoneType" in caplog.text
|
|
|
|
@patch("aareDB.ApiClient")
|
|
def test_format_scan_payload_error(mock_api, mock_bl, sample_info, caplog):
|
|
wrapper = AareWrapper(bl=mock_bl)
|
|
with pytest.raises(Exception):
|
|
# Passing None for geom should trigger error when accessing geom.beam_size_mm
|
|
wrapper.format_scan_payload(sample_info, None, None, (0,0))
|
|
assert "NoneType" in caplog.text
|