AareDAQ/tests/unit/daq/operations/mounting/test_mounting_service.py

import types
import sys

from aare.common.coordinate import AerotechCoordinate, Coordinate
from aare.common.exception_handler import CriticalTellException, DoorSafetyError, MountingFailed
from aare.common.models import DewarAddress, SampleShortInfo
from aare.devices.tell_client import TellEventValueEnum
from aare.daq.operations.mounting.models import MountingContext, MountingResult, MountingDependencies, MountingSettings
from aare.daq.operations.mounting.service import MountingService


if "jfjoch_client.models.scan_result" not in sys.modules:
    jfjoch_client_mod = types.ModuleType("jfjoch_client")
    jfjoch_client_models_mod = types.ModuleType("jfjoch_client.models")
    jfjoch_client_scan_result_mod = types.ModuleType("jfjoch_client.models.scan_result")

    jfjoch_client_scan_result_mod.ScanResult = dict
    jfjoch_client_models_mod.scan_result = jfjoch_client_scan_result_mod
    jfjoch_client_mod.models = jfjoch_client_models_mod
    sys.modules["jfjoch_client"] = jfjoch_client_mod
    sys.modules["jfjoch_client.models"] = jfjoch_client_models_mod
    sys.modules["jfjoch_client.models.scan_result"] = jfjoch_client_scan_result_mod


def _make_sample(sample_id: int, name: str) -> SampleShortInfo:
    return SampleShortInfo(
        db_id=sample_id,
        puck_name="puck1",
        dewar_name="dew1",
        sample_name=name,
        run_number=1,
        user="group1",
        pin=sample_id,
        location=DewarAddress(segment="A", pos=1),
    )


def _make_context(previous_sample=None, *, prohibited=True, alarm=False):
    tell = types.SimpleNamespace(
        mount=lambda **kwargs: TellEventValueEnum.SUCCESS,
        unmount=lambda **kwargs: None,
        dry=lambda **kwargs: None,
        check_enable_motion=lambda: None,
        wait_not_busy=lambda timeout=360.0: None,
        validate_door_closed=lambda: None,
        set_in_mount_position=lambda value: None,
    )

    devs = types.SimpleNamespace(
        smargon_move_home=lambda: None,
        aerotech_pos=None,
        tell=tell,
        magnet_position_sensor=types.SimpleNamespace(value=0),
        pss=types.SimpleNamespace(
            is_prohibited=lambda: prohibited,
            alarm_active=lambda: alarm,
        ),
    )

    streak = {"count": 0}

    cfg = types.SimpleNamespace(
        current_sample=previous_sample,
        get_mount_failure_streak=lambda: streak["count"],
        increment_mount_failure_streak=lambda: streak.__setitem__("count", streak["count"] + 1) or streak["count"],
        reset_mount_failure_streak=lambda: streak.__setitem__("count", 0),
        record_mount_failure=lambda: streak.__setitem__("count", streak["count"] + 1) or streak["count"],
        record_mount_success=lambda: streak.__setitem__("count", 0),
    )

    return MountingContext(
        deps=MountingDependencies(
            cfg=cfg,
            devs=devs,
        ),
        settings=MountingSettings(
            mount_position=AerotechCoordinate(at_mm=Coordinate(x=0, y=0, z=0), omega_deg=0),
        ),
    )


def test_execute_mount_success(mock_logger):
    previous_sample = _make_sample(1, "old")
    target_sample = _make_sample(2, "new")
    ctx = _make_context(previous_sample=previous_sample)

    service = MountingService(context=ctx, logger=mock_logger)

    result = service.execute(target=target_sample)

    assert isinstance(result, MountingResult)
    assert result.success is True
    assert result.mounted_sample == target_sample
    assert result.previous_sample == previous_sample
    assert result.did_unmount_previous is True
    assert ctx.deps.cfg.current_sample == target_sample


def test_execute_unmount_success(mock_logger):
    previous_sample = _make_sample(1, "old")
    ctx = _make_context(previous_sample=previous_sample)

    service = MountingService(context=ctx, logger=mock_logger)

    result = service.execute(target=None)

    assert result.success is True
    assert result.mounted_sample is None
    assert result.previous_sample == previous_sample
    assert result.did_unmount_previous is True
    assert ctx.deps.cfg.current_sample is None


def test_execute_mount_returns_failed_result_for_no_pin_in_gripper(mock_logger):
    target_sample = _make_sample(2, "new")
    ctx = _make_context()
    ctx.deps.devs.tell.mount = lambda **kwargs: TellEventValueEnum.NO_PIN_IN_GRIPPER

    service = MountingService(context=ctx, logger=mock_logger)

    result = service.execute(target=target_sample)

    assert result.success is False
    assert isinstance(result.error, MountingFailed)
    assert result.is_error is True


def test_execute_mount_returns_failed_result_for_unhandled_tell_response(mock_logger):
    target_sample = _make_sample(2, "new")
    ctx = _make_context()
    ctx.deps.devs.tell.mount = lambda **kwargs: TellEventValueEnum.UNKNOWN

    service = MountingService(context=ctx, logger=mock_logger)

    result = service.execute(target=target_sample)

    assert result.success is False
    assert isinstance(result.error, CriticalTellException)
    assert result.is_error is True


def test_dry_unmounts_current_sample_before_drying(mock_logger):
    previous_sample = _make_sample(1, "old")
    ctx = _make_context(previous_sample=previous_sample)
    unmount_calls = []
    dry_calls = []

    ctx.deps.devs.tell.unmount = lambda **kwargs: unmount_calls.append(kwargs)
    ctx.deps.devs.tell.dry = lambda **kwargs: dry_calls.append(kwargs)

    service = MountingService(context=ctx, logger=mock_logger)
    service.dry(park=True, unmount=True)

    assert len(unmount_calls) == 1
    assert ctx.deps.cfg.current_sample is None
    assert dry_calls == [{"wait_cold": -1, "wait": True}]


def test_execute_mount_triggers_dry_on_third_consecutive_failure(mock_logger):
    target_sample = _make_sample(2, "new")
    dry_calls = []

    ctx = _make_context()
    ctx.deps.devs.tell.mount = lambda **kwargs: TellEventValueEnum.NO_PIN_IN_GRIPPER
    ctx.deps.devs.tell.dry = lambda **kwargs: dry_calls.append(kwargs)
    ctx.deps.cfg.increment_mount_failure_streak()
    ctx.deps.cfg.increment_mount_failure_streak()

    service = MountingService(context=ctx, logger=mock_logger)

    result = service.execute(target=target_sample)

    assert result.success is False
    assert isinstance(result.error, MountingFailed)
    assert result.error.critical is False
    assert dry_calls == [{"wait": True}]


def test_execute_mount_stops_automation_on_fifth_consecutive_failure(mock_logger):
    target_sample = _make_sample(2, "new")
    dry_calls = []

    ctx = _make_context()
    ctx.deps.devs.tell.mount = lambda **kwargs: TellEventValueEnum.NO_PIN_IN_GRIPPER
    ctx.deps.devs.tell.dry = lambda **kwargs: dry_calls.append(kwargs)
    ctx.deps.cfg.increment_mount_failure_streak()
    ctx.deps.cfg.increment_mount_failure_streak()
    ctx.deps.cfg.increment_mount_failure_streak()
    ctx.deps.cfg.increment_mount_failure_streak()

    service = MountingService(context=ctx, logger=mock_logger)

    result = service.execute(target=target_sample)

    assert result.success is False
    assert isinstance(result.error, MountingFailed)
    assert result.error.critical is True
    assert "Mount failed 5 times in a row" in str(result.error)
    assert {"wait_cold": -1, "wait": True} in dry_calls


def test_execute_mount_success_resets_failure_streak(mock_logger):
    previous_sample = _make_sample(1, "old")
    target_sample = _make_sample(2, "new")
    ctx = _make_context(previous_sample=previous_sample)
    ctx.deps.cfg.increment_mount_failure_streak()
    ctx.deps.cfg.increment_mount_failure_streak()

    service = MountingService(context=ctx, logger=mock_logger)

    result = service.execute(target=target_sample)

    assert result.success is True
    assert ctx.deps.cfg.get_mount_failure_streak() == 0


def test_execute_mount_critical_tell_error_does_not_increment_failure_streak(mock_logger):
    target_sample = _make_sample(2, "new")
    ctx = _make_context()

    def mount(**kwargs):
        raise CriticalTellException("critical tell problem")

    ctx.deps.devs.tell.mount = mount

    service = MountingService(context=ctx, logger=mock_logger)

    result = service.execute(target=target_sample)

    assert result.success is False
    assert isinstance(result.error, CriticalTellException)
    assert ctx.deps.cfg.get_mount_failure_streak() == 0

def test_mount_blocked_when_not_prohibited(mock_logger):
    """Doors open / hutch not in prohibited state -> critical DoorSafetyError,
    and it is not counted as a mount-failure streak."""
    target_sample = _make_sample(2, "new")
    ctx = _make_context(prohibited=False)

    service = MountingService(context=ctx, logger=mock_logger)

    result = service.execute(target=target_sample)

    assert result.success is False
    assert isinstance(result.error, DoorSafetyError)
    assert result.error.critical is True
    # Door safety is a pre-flight gate, not a mount attempt: streak stays clear.
    assert ctx.deps.cfg.get_mount_failure_streak() == 0


def test_mount_blocked_when_alarm_active(mock_logger):
    target_sample = _make_sample(3, "new")
    ctx = _make_context(alarm=True)

    service = MountingService(context=ctx, logger=mock_logger)

    result = service.execute(target=target_sample)

    assert result.success is False
    assert isinstance(result.error, DoorSafetyError)
    assert result.error.critical is True


def test_unmount_blocked_when_not_prohibited_keeps_door_error(mock_logger):
    """Unmount path must not swallow DoorSafetyError into UnmountingFailed."""
    previous_sample = _make_sample(1, "old")
    ctx = _make_context(previous_sample=previous_sample, prohibited=False)

    service = MountingService(context=ctx, logger=mock_logger)

    result = service.execute(target=None)

    assert result.success is False
    assert isinstance(result.error, DoorSafetyError)
    assert result.error.critical is True