import sys import types from aarecommon.errors.exception_handler import ( CriticalTellException, DoorSafetyError, MountingFailed, ) from aarecommon.math.coordinate import AerotechCoordinate, Coordinate from aarecommon.models.models import DewarAddress, SampleShortInfo from aare.daq.operations.mounting.models import ( MountingContext, MountingDependencies, MountingResult, MountingSettings, ) from aare.daq.operations.mounting.service import MountingService from aare.devices.tell_client import TellEventValueEnum 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