import types from unittest.mock import MagicMock, patch from aarecommon.errors.exception_handler import LoopCenteringFailed from aarecommon.models.automation import ( AutomationProgress, StepState, StepStatus, WorkflowStateKind, ) from aarecommon.models.models import SampleShortInfo from aare.daq.daq import AareDAQ def test_execute_loop_centering_rejects_invalid_sample(): daq = object.__new__(AareDAQ) daq._AareDAQ__set_state = MagicMock() daq._AareDAQ__devs = MagicMock() daq._AareDAQ__aare = MagicMock() daq._handle_operation_error = MagicMock() daq._create_loop_centering_service = MagicMock() result = daq._execute_loop_centering(None) assert result is False daq._create_loop_centering_service.assert_not_called() def test_execute_loop_centering_handles_service_failure(): daq = object.__new__(AareDAQ) daq._AareDAQ__set_state = MagicMock() daq._AareDAQ__devs = types.SimpleNamespace(lamp_light=None) daq._AareDAQ__aare = MagicMock() daq._handle_operation_error = MagicMock() daq.save_screenshot_db = MagicMock() sample = SampleShortInfo( db_id=1, sample_name="S1", puck_name="P1", dewar_name="D1", run_number=1, user="p1", pin=1, location=None, ) service = MagicMock() service.run.return_value = types.SimpleNamespace( success=False, comment="failed", error=LoopCenteringFailed("failed") ) daq._create_loop_centering_service = MagicMock(return_value=service) result = daq._execute_loop_centering(sample) assert result is False daq._handle_operation_error.assert_called_once() def test_execute_loop_centering_handles_service_success(): daq = object.__new__(AareDAQ) daq._AareDAQ__set_state = MagicMock() daq._AareDAQ__devs = MagicMock() daq._AareDAQ__aare = MagicMock() daq._handle_operation_error = MagicMock() daq.save_screenshot_db = MagicMock() sample = SampleShortInfo( db_id=1, sample_name="S1", puck_name="P1", dewar_name="D1", run_number=1, user="p1", pin=1, location=None, ) service = MagicMock() service.run.return_value = types.SimpleNamespace(success=True, comment=None, error=None) daq._create_loop_centering_service = MagicMock(return_value=service) result = daq._execute_loop_centering(sample) assert result is True daq.save_screenshot_db.assert_called_once() def test_record_best_effort_step_failure_marks_progress_and_logs_warning(mock_logger): daq = object.__new__(AareDAQ) daq._automation_progress_cb = None daq._set_progress_step = AareDAQ._set_progress_step.__get__(daq, AareDAQ) progress = AutomationProgress( current_step="Center", steps=[ StepState( step=WorkflowStateKind.LOOP_CENTRE, status=StepStatus.RUNNING, message="Centering" ) ], finished=False, success=None, ) sample = SampleShortInfo( db_id=1, sample_name="S1", puck_name="P1", dewar_name="D1", run_number=1, user="p1", pin=1, location=None, ) with patch("aare.daq.daq.logger", mock_logger): daq._record_best_effort_step_failure( progress=progress, step=WorkflowStateKind.LOOP_CENTRE, error=LoopCenteringFailed("Loop centering failed"), sample=sample, code="LOOP_CENTERING_FAILED", ) step = progress.steps[0] assert step.status == StepStatus.FAILED assert step.message == "Loop centering failed" assert len(progress.events) == 1 assert progress.events[0].level == "WARNING" assert progress.events[0].code == "LOOP_CENTERING_FAILED" assert progress.events[0].exception_class == "LoopCenteringFailed" assert progress.events[0].message == "Loop centering failed" assert progress.events[0].sample_id == 1 mock_logger.warning.assert_called_once()