AareDAQ/tests/unit/daq/test_aare_daq_loop_centering.py

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(mock_daq):
    mock_daq._create_loop_centering_service = MagicMock()

    result = mock_daq._execute_loop_centering(None)

    assert result is False
    mock_daq._create_loop_centering_service.assert_not_called()


def test_execute_loop_centering_handles_service_failure(mock_daq):
    mock_daq._devs = types.SimpleNamespace(lamp_light=None)
    mock_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")
    )
    mock_daq._create_loop_centering_service = MagicMock(return_value=service)

    result = mock_daq._execute_loop_centering(sample)

    assert result is False
    mock_daq._handle_operation_error.assert_called_once()


def test_execute_loop_centering_handles_service_success(mock_daq):
    mock_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)
    mock_daq._create_loop_centering_service = MagicMock(return_value=service)

    result = mock_daq._execute_loop_centering(sample)

    assert result is True
    mock_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()