diff --git a/tests/unit/daq/test_aaredb.py b/tests/unit/daq/test_aaredb.py new file mode 100644 index 00000000..c93d4511 --- /dev/null +++ b/tests/unit/daq/test_aaredb.py @@ -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