tests: added aaredb wrapper test
This commit is contained in:
@@ -0,0 +1,322 @@
|
||||
import pytest
|
||||
import numpy as np
|
||||
import io
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch, ANY
|
||||
from aare.daq.aaredb import AareWrapper, time_db_call
|
||||
from aare.common.beamline import MXBeamline
|
||||
from aare.common.models import (
|
||||
SampleShortInfo, PuckLoadedInfo, DewarAddress,
|
||||
DAQStatusModel, ScanResultPayloadModel
|
||||
)
|
||||
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)
|
||||
)
|
||||
|
||||
def test_time_db_call_decorator_logs_duration(caplog):
|
||||
import logging
|
||||
caplog.set_level(logging.DEBUG)
|
||||
|
||||
@time_db_call
|
||||
def slow_func():
|
||||
import time
|
||||
time.sleep(0.01)
|
||||
return "done"
|
||||
|
||||
result = slow_func()
|
||||
assert result == "done"
|
||||
assert "Starting AareDB call: slow_func" in caplog.text
|
||||
assert "Finished AareDB call: slow_func in" in caplog.text
|
||||
|
||||
def test_time_db_call_decorator_logs_failure(caplog):
|
||||
import logging
|
||||
caplog.set_level(logging.DEBUG)
|
||||
|
||||
@time_db_call
|
||||
def fail_func():
|
||||
raise ValueError("oops")
|
||||
|
||||
with pytest.raises(ValueError, match="oops"):
|
||||
fail_func()
|
||||
|
||||
assert "Starting AareDB call: fail_func" in caplog.text
|
||||
assert "fail_func FAILED after" in caplog.text
|
||||
assert "with error: oops" in caplog.text
|
||||
|
||||
@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
|
||||
Reference in New Issue
Block a user