87 lines
3.5 KiB
Python
87 lines
3.5 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
from aare.devices.aerotech import AerotechController, AEROTECH_HOME
|
|
from aare.common.beamline import MXBeamline
|
|
from aare.common.coordinate import AerotechCoordinate, Coordinate
|
|
from aare.common.exception_handler import AerotechCommunicationError
|
|
|
|
@pytest.fixture
|
|
def mock_aerotech_api():
|
|
with patch('aarescan_client.ApiClient'), \
|
|
patch('aarescan_client.DefaultApi') as mock_api_class, \
|
|
patch('aarescan_client.Configuration'):
|
|
mock_api = mock_api_class.return_value
|
|
yield mock_api
|
|
|
|
@pytest.fixture
|
|
def aerotech_controller(mock_aerotech_api):
|
|
# Set environment variable before initialization
|
|
with patch.dict("os.environ", {"BEAMLINE": "X10SA"}):
|
|
controller = AerotechController(MXBeamline.X10SA)
|
|
|
|
# Manually inject the mocked API into the controller's private attribute
|
|
# This ensures that calls to self.__api go to our mock.
|
|
controller._AerotechController__api = mock_aerotech_api
|
|
controller._AerotechController__simulated = False
|
|
return controller
|
|
|
|
def test_init_x10sa(mock_aerotech_api):
|
|
controller = AerotechController(MXBeamline.X10SA)
|
|
assert controller._AerotechController__base == "http://mx-x10sa-queue-01.psi.ch:5234"
|
|
assert controller._AerotechController__simulated is False
|
|
|
|
def test_init_simulated():
|
|
controller = AerotechController(MXBeamline.SIMULATED)
|
|
assert controller._AerotechController__simulated is True
|
|
|
|
def test_cancel(aerotech_controller, mock_aerotech_api):
|
|
aerotech_controller.cancel()
|
|
mock_aerotech_api.cancel_post.assert_called_once()
|
|
|
|
def test_is_idle(aerotech_controller, mock_aerotech_api):
|
|
mock_aerotech_api.status_get.return_value.state = 'Idle'
|
|
assert aerotech_controller.is_idle() is True
|
|
|
|
mock_aerotech_api.status_get.return_value.state = 'Busy'
|
|
assert aerotech_controller.is_idle() is False
|
|
|
|
def test_get_position(aerotech_controller, mock_aerotech_api):
|
|
status = MagicMock()
|
|
status.x.pos = 1.0
|
|
status.y.pos = 2.0
|
|
status.z.pos = 3.0
|
|
status.u.pos = 90.0
|
|
mock_aerotech_api.status_get.return_value = status
|
|
|
|
pos = aerotech_controller.get_position()
|
|
assert pos.at_mm.x == 1.0
|
|
assert pos.at_mm.y == 2.0
|
|
assert pos.at_mm.z == 3.0
|
|
assert pos.omega_deg == 90.0
|
|
|
|
def test_move_home(aerotech_controller, mock_aerotech_api):
|
|
aerotech_controller.move_home()
|
|
mock_aerotech_api.position_post.assert_called_once()
|
|
|
|
def test_position(aerotech_controller, mock_aerotech_api):
|
|
target = AerotechCoordinate(at_mm=Coordinate(x=1, y=2, z=3), omega_deg=45)
|
|
aerotech_controller.position(target)
|
|
mock_aerotech_api.position_post.assert_called_once()
|
|
|
|
def test_rotation_scan(aerotech_controller, mock_aerotech_api):
|
|
aerotech_controller.rotation_scan(rotation_deg=360, time_sec=10, start_pos_deg=0)
|
|
mock_aerotech_api.rotation_scan_post.assert_called_once()
|
|
|
|
def test_grid_scan(aerotech_controller, mock_aerotech_api):
|
|
aerotech_controller.grid_scan(grid_elem_count_y=10, grid_elem_size_y_um=10, grid_elem_count_x=10, grid_elem_size_x_um=10, time_sec=5)
|
|
mock_aerotech_api.grid_scan_post.assert_called_once()
|
|
|
|
def test_screening_scan(aerotech_controller, mock_aerotech_api):
|
|
aerotech_controller.screening_scan(rotation_deg=10, wedge_deg=2, time_sec=1, steps=5)
|
|
mock_aerotech_api.screening_post.assert_called_once()
|
|
|
|
def test_api_error(aerotech_controller, mock_aerotech_api):
|
|
mock_aerotech_api.status_get.side_effect = Exception("API Error")
|
|
with pytest.raises(AerotechCommunicationError):
|
|
aerotech_controller.is_idle()
|