diff --git a/tests/unit/common/test_aare_exception.py b/tests/unit/common/test_aare_exception.py new file mode 100644 index 00000000..909cdefc --- /dev/null +++ b/tests/unit/common/test_aare_exception.py @@ -0,0 +1,199 @@ +"""Tests for the AareException hierarchy introduced in Phase 1 of the +exception-handling redesign. + +These cover: +- class-level ``critical`` default +- instance-level override via ``critical=...`` kwarg +- correct parent/family relationships so watchers can match by family +""" + +from __future__ import annotations + +import pytest + +from aare.common.exception_handler import ( + AareException, + AutomationError, + AareUserError, + AareAuthError, + TellException, + SmargonException, + AerotechException, + JFJochException, + BECException, + BeamlineStateException, + DataCollectionException, + RasterScanException, + AutoRasterSampleSkipped, + LoopCenteringFailed, + MountingFailed, + UnmountingFailed, + TellCommunicationError, + TellConnectionException, + TellCommandWhileBusyException, + TellMountFailedException, + WarningTellException, + CriticalTellException, + SmargonCommunicationError, + AerotechCommunicationError, + JFJochCommunicationError, + BECCommunicationError, + AareDBCommunicationError, + MagnetPositionSensorErorr, + SmartMagnetFaultException, + TransformationInvalidException, + StateTransitionFailed, + MaintenanceStateException, + BeamlineBusyException, + BeamlineBusyTimeoutException, + AXCFailed, + SampleException, + ManualMountException, + AuthenticationException, + UserRightsException, +) + + +# --------------------------------------------------------------------------- +# Class-level criticality defaults +# --------------------------------------------------------------------------- + +def test_aare_exception_class_critical_default_false(): + assert AareException.critical is False + + +def test_subclasses_inherit_class_critical_default(): + # All Phase 1 classes inherit critical=False until Phase 3 sets explicit flags + for cls in (AutomationError, AareUserError, AareAuthError, + TellException, SmargonException, AerotechException, + JFJochException, BECException, BeamlineStateException, + MountingFailed, UnmountingFailed, LoopCenteringFailed, + TellCommunicationError, AareDBCommunicationError): + assert cls.critical is False, f"{cls.__name__} should default to critical=False in Phase 1" + + +# --------------------------------------------------------------------------- +# Instance-level critical override +# --------------------------------------------------------------------------- + +def test_instance_critical_override_true(): + exc = MountingFailed("mount fail", critical=True) + assert exc.critical is True + # class default unchanged + assert MountingFailed.critical is False + + +def test_instance_critical_override_false(): + exc = MountingFailed("mount fail", critical=False) + assert exc.critical is False + + +def test_instance_critical_none_keeps_class_default(): + # Omitting critical uses class default + exc = MountingFailed("mount fail") + assert exc.critical is False # class default + + +def test_aaredb_communication_error_instance_critical_override(): + exc = AareDBCommunicationError("db down", critical=True, operation="GET") + assert exc.critical is True + assert exc.operation == "GET" + + +# --------------------------------------------------------------------------- +# Routing-root membership +# --------------------------------------------------------------------------- + +def test_automation_errors_are_automation_error(): + for cls in (LoopCenteringFailed, MountingFailed, UnmountingFailed, + AXCFailed, TellCommunicationError, SmargonCommunicationError, + AerotechCommunicationError, JFJochCommunicationError, + BECCommunicationError, AareDBCommunicationError, + BeamlineBusyException, RasterScanException, + MagnetPositionSensorErorr, SmartMagnetFaultException, + TransformationInvalidException, AutoRasterSampleSkipped): + exc = cls() if cls is not AutoRasterSampleSkipped else cls("skipped") + assert isinstance(exc, AutomationError), f"{cls.__name__} should be AutomationError" + assert isinstance(exc, AareException), f"{cls.__name__} should be AareException" + + +def test_auth_errors_are_aare_auth_error(): + for cls in (AuthenticationException, UserRightsException): + exc = cls() + assert isinstance(exc, AareAuthError) + assert isinstance(exc, AareException) + + +def test_user_errors_are_aare_user_error(): + for cls in (ManualMountException, SampleException): + exc = cls() + assert isinstance(exc, AareUserError) + assert isinstance(exc, AareException) + + +# --------------------------------------------------------------------------- +# Family-base membership (drives watcher matching) +# --------------------------------------------------------------------------- + +def test_tell_family_membership(): + for cls in (TellCommunicationError, TellConnectionException, + CriticalTellException, TellCommandWhileBusyException, + WarningTellException, MountingFailed, UnmountingFailed, + TellMountFailedException): + exc = cls() + assert isinstance(exc, TellException), f"{cls.__name__} should be in Tell family" + assert isinstance(exc, AutomationError) + + +def test_smargon_family_membership(): + exc = SmargonCommunicationError() + assert isinstance(exc, SmargonException) + assert isinstance(exc, AutomationError) + + +def test_aerotech_family_membership(): + exc = AerotechCommunicationError() + assert isinstance(exc, AerotechException) + assert isinstance(exc, AutomationError) + + +def test_jfjoch_family_membership(): + exc = JFJochCommunicationError() + assert isinstance(exc, JFJochException) + assert isinstance(exc, AutomationError) + + +def test_bec_family_membership(): + exc = BECCommunicationError() + assert isinstance(exc, BECException) + assert isinstance(exc, AutomationError) + + +def test_beamline_state_family_membership(): + for cls in (StateTransitionFailed, MaintenanceStateException, + BeamlineBusyException, BeamlineBusyTimeoutException): + exc = cls() + assert isinstance(exc, BeamlineStateException) + assert isinstance(exc, AutomationError) + + +def test_data_collection_family_membership(): + exc = RasterScanException() + assert isinstance(exc, DataCollectionException) + assert isinstance(exc, AutomationError) + + +# --------------------------------------------------------------------------- +# Backward compat: existing __str__/.message contracts preserved +# --------------------------------------------------------------------------- + +def test_message_attribute_preserved(): + exc = LoopCenteringFailed("custom message") + assert exc.message == "custom message" + assert str(exc) == "custom message" + + +def test_default_messages_preserved(): + assert str(MountingFailed()) == "A sample was not mounted" + assert str(UnmountingFailed()) == "A sample was not unmounted" + assert str(LoopCenteringFailed()) == "Loop Centering did not detect a sample" diff --git a/tests/unit/daq/operations/mounting/test_mounting_service.py b/tests/unit/daq/operations/mounting/test_mounting_service.py new file mode 100644 index 00000000..e4ca825e --- /dev/null +++ b/tests/unit/daq/operations/mounting/test_mounting_service.py @@ -0,0 +1,129 @@ +import types + +from aare.common.coordinate import AerotechCoordinate, Coordinate +from aare.common.exception_handler import CriticalTellException, MountingFailed +from aare.common.models import DewarAddress, SampleShortInfo +from aare.devices.tell_client import TellEventValueEnum +from aare.daq.operations.mounting.models import MountingContext, MountingResult +from aare.daq.operations.mounting.service import MountingService + + +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): + 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, + 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), + ) + + cfg = types.SimpleNamespace( + current_sample=previous_sample, + get_mount_fail_count=lambda: 0, + record_mount_success=lambda: None, + record_mount_failure=lambda: 1, + ) + + return MountingContext( + cfg=cfg, + devs=devs, + 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.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.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.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.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.devs.tell.unmount = lambda **kwargs: unmount_calls.append(kwargs) + ctx.devs.tell.dry = lambda **kwargs: dry_calls.append(kwargs) + + service = MountingService(context=ctx, logger=mock_logger) + service.dry(park=True) + + assert len(unmount_calls) == 1 + assert ctx.cfg.current_sample is None + assert dry_calls == [{"wait_cold": -1, "wait": True}] diff --git a/tests/unit/daq/operations/screenshot/test_screenshot_service.py b/tests/unit/daq/operations/screenshot/test_screenshot_service.py new file mode 100644 index 00000000..5680da8e --- /dev/null +++ b/tests/unit/daq/operations/screenshot/test_screenshot_service.py @@ -0,0 +1,105 @@ +import types + +import numpy as np +import pytest + +from aare.common.models import DewarAddress, SampleShortInfo +from aare.daq.operations.screenshot.service import ScreenshotService + + +def _make_sample(sample_id: int = 7) -> SampleShortInfo: + return SampleShortInfo( + db_id=sample_id, + puck_name="puck1", + dewar_name="dew1", + sample_name="sample1", + run_number=1, + user="p12345", + pin=1, + location=DewarAddress(segment="A", pos=1), + ) + + +def test_save_to_db_uploads_image_via_shared_aare(mock_logger): + image = np.zeros((10, 10, 3), dtype=np.uint8) + upload_calls = [] + noncritical_calls = [] + + mlbox = types.SimpleNamespace(get_latest_image=lambda: image) + aare = types.SimpleNamespace( + upload_image=lambda sample_id, filename, bgr_image, **kwargs: upload_calls.append( + (sample_id, filename, bgr_image, kwargs) + ) + ) + + sample = _make_sample(11) + + def _run_noncritical(action, *, description, sample=None): + noncritical_calls.append((description, sample)) + return action() + + service = ScreenshotService( + mlbox=mlbox, + aare=aare, + logger=mock_logger, + run_noncritical=_run_noncritical, + sample_provider=lambda: sample, + pgroup_provider=lambda: "p12345", + ) + + service.save_to_db(11, "mounted", 0.0) + + assert len(upload_calls) == 1 + assert upload_calls[0][0] == 11 + assert upload_calls[0][1] == "mounted" + assert noncritical_calls == [("screenshot upload 'mounted'", sample)] + + +def test_send_to_db_writes_photo_and_uploads_with_default_message(mock_logger, tmp_path): + image = np.zeros((10, 10, 3), dtype=np.uint8) + upload_calls = [] + + mlbox = types.SimpleNamespace(get_latest_image=lambda: image) + aare = types.SimpleNamespace( + upload_image=lambda sample_id, filename, bgr_image, **kwargs: upload_calls.append( + (sample_id, filename, kwargs) + ) + ) + sample = _make_sample(15) + + service = ScreenshotService( + mlbox=mlbox, + aare=aare, + logger=mock_logger, + run_noncritical=lambda action, **kwargs: action(), + sample_provider=lambda: sample, + pgroup_provider=lambda: "p12345", + photos_root=str(tmp_path), + ) + + service.send_to_db( + filename="test image", + message=None, + default_message="default screenshot message", + ) + + expected_photo = tmp_path / "p12345" / "raw" / "photos" / "15" / "test_image.jpeg" + assert expected_photo.exists() + assert len(upload_calls) == 1 + assert upload_calls[0][0] == 15 + assert upload_calls[0][1] == "test_image" + assert upload_calls[0][2]["message"] == "default screenshot message" + + +def test_send_to_db_requires_mounted_sample(mock_logger): + service = ScreenshotService( + mlbox=types.SimpleNamespace(get_latest_image=lambda: np.zeros((4, 4, 3), dtype=np.uint8)), + aare=types.SimpleNamespace(upload_image=lambda *args, **kwargs: None), + logger=mock_logger, + run_noncritical=lambda action, **kwargs: action(), + sample_provider=lambda: None, + pgroup_provider=lambda: "p12345", + ) + + with pytest.raises(ValueError, match="valid sample_id"): + service.send_to_db(default_message="x") \ No newline at end of file diff --git a/tests/unit/daq/test_gui_timeout.py b/tests/unit/daq/test_gui_timeout.py new file mode 100644 index 00000000..31015276 --- /dev/null +++ b/tests/unit/daq/test_gui_timeout.py @@ -0,0 +1,203 @@ +from unittest.mock import patch + +from starlette.testclient import TestClient + +from aare.daq.server import app + + +def test_status_renews_gui_session_with_gui_timeout(client, mock_backend): + mock_cfg = mock_backend["cfg"] + mock_daq = mock_backend["daq"] + + from aare.common.models import ( + DAQStatusModel, + SessionStatus, + SessionsStateEnum, + BeamlineStateEnum, + SampleGeometryModel, + BeamlineStatus, + CrystalSize, + SampleCameraSettings, + ) + from aare.common.diffraction_geometry import DiffractionGeometry + from aare.common.coordinate import Coordinate, SmargonCoordinate + + geom = SampleGeometryModel( + beam_location_pxl=Coordinate(x=500, y=500), + pixel_in_mm=0.001, + aerotech=Coordinate(x=0, y=0, z=0), + aerotech_meas=Coordinate(x=0, y=0, z=0), + smargon=SmargonCoordinate(sh_mm=Coordinate(x=0, y=0, z=0), phi_deg=0.0, chi_deg=0.0), + omega_deg=0.0, + beam_size_mm=Coordinate(x=0.01, y=0.01), + ) + diff = DiffractionGeometry( + energy_keV=12.0, + dtz_mm=150.0, + pixel_size_mm=0.075, + beam_center_pxl=(1000.0, 1000.0), + detector_size_pxl=(2000, 2000), + detector_description="Eiger 16M", + detector_serial_number="123", + poni_rot1_rad=0.0, + poni_rot2_rad=0.0, + ) + bl_status = BeamlineStatus( + name="X06DA", + ring_current_mA=400.0, + front_light=0.0, + back_light=0.0, + cryojet_K=100.0, + shutter_open=False, + exp_shutter_open=False, + flux_ph_s=1e12, + transmission=1.0, + zoom=1.0, + sample_camera=SampleCameraSettings(exposure=0.1, gain=1.0), + commissioning_mode=False, + dtz_min=120.0, + dtz_max=1600.0, + ) + session_status = SessionStatus( + session=SessionsStateEnum.Vacant, + current_pgroup="p12345", + staff=True, + ) + + mock_daq.status = DAQStatusModel( + geom=geom, + diffraction=diff, + bl=bl_status, + state=BeamlineStateEnum.Maintenance, + busy=False, + session=session_status, + crystal_size=CrystalSize(x=0, y=0, z=0), + ) + mock_cfg.pgroup = "p12345" + mock_cfg.session_state.return_value = SessionsStateEnum.OwnedByYou + mock_cfg.get_open_gui_sessions.return_value = [] + mock_cfg.GUI_SESSION_EXPIRE_SECONDS = 10 + + with patch("aare.daq.server.cfg", mock_cfg), patch("aare.daq.server.daq", mock_daq): + response = client.get("/status", headers={"Authorization": "Bearer fake-token"}) + + assert response.status_code == 200 + mock_cfg.touch_gui_session.assert_called_once_with( + session=123, + username="testuser", + staff=True, + expiry_sec=10, + ) + + +def test_status_hides_open_guis_for_non_staff(mock_backend): + with patch("aare.daq.auth.parse_token") as mock_parse: + from aare.daq.auth import TokenData + from aare.common.models import ( + DAQStatusModel, + SessionStatus, + SessionsStateEnum, + BeamlineStateEnum, + SampleGeometryModel, + BeamlineStatus, + CrystalSize, + SampleCameraSettings, + OpenGuiSessionInfo, + ) + from aare.common.diffraction_geometry import DiffractionGeometry + from aare.common.coordinate import Coordinate, SmargonCoordinate + + mock_parse.return_value = TokenData( + sub="user1", + staff=False, + pgroups=["p12345"], + session=123, + ) + + geom = SampleGeometryModel( + beam_location_pxl=Coordinate(x=500, y=500), + pixel_in_mm=0.001, + aerotech=Coordinate(x=0, y=0, z=0), + aerotech_meas=Coordinate(x=0, y=0, z=0), + smargon=SmargonCoordinate(sh_mm=Coordinate(x=0, y=0, z=0), phi_deg=0.0, chi_deg=0.0), + omega_deg=0.0, + beam_size_mm=Coordinate(x=0.01, y=0.01), + ) + diff = DiffractionGeometry( + energy_keV=12.0, + dtz_mm=150.0, + pixel_size_mm=0.075, + beam_center_pxl=(1000.0, 1000.0), + detector_size_pxl=(2000, 2000), + detector_description="Eiger 16M", + detector_serial_number="123", + poni_rot1_rad=0.0, + poni_rot2_rad=0.0, + ) + bl_status = BeamlineStatus( + name="X06DA", + ring_current_mA=400.0, + front_light=0.0, + back_light=0.0, + cryojet_K=100.0, + shutter_open=False, + exp_shutter_open=False, + flux_ph_s=1e12, + transmission=1.0, + zoom=1.0, + sample_camera=SampleCameraSettings(exposure=0.1, gain=1.0), + commissioning_mode=False, + dtz_min=120.0, + dtz_max=1600.0, + ) + + mock_backend["daq"].status = DAQStatusModel( + geom=geom, + diffraction=diff, + bl=bl_status, + state=BeamlineStateEnum.Maintenance, + busy=False, + session=SessionStatus(session=SessionsStateEnum.Vacant, current_pgroup="p12345", staff=False), + crystal_size=CrystalSize(x=0, y=0, z=0), + ) + mock_backend["cfg"].pgroup = "p12345" + mock_backend["cfg"].session_state.return_value = SessionsStateEnum.OwnedByElse + mock_backend["cfg"].get_open_gui_sessions.return_value = [ + OpenGuiSessionInfo(session=1, username="staff1", last_seen_ts=1.0, staff=True) + ] + mock_backend["cfg"].get_gui_session.return_value = OpenGuiSessionInfo( + session=123, + username="user1", + last_seen_ts=2.0, + staff=False, + ) + + with TestClient(app) as local_client: + response = local_client.get("/status", headers={"Authorization": "Bearer fake-token"}) + + assert response.status_code == 200 + payload = response.json()["open_guis"] + assert len(payload) == 1 + assert payload[0]["session"] == 123 + assert payload[0]["username"] == "user1" + + +def test_admin_gui_sessions_includes_baton_holder_flag(client, mock_backend): + from aare.common.models import OpenGuiSessionInfo + + mock_cfg = mock_backend["cfg"] + mock_cfg.get_open_gui_sessions.return_value = [ + OpenGuiSessionInfo(session=11, username="user1", last_seen_ts=1.0, staff=False, holds_baton=False), + OpenGuiSessionInfo(session=22, username="holder", last_seen_ts=2.0, staff=True, holds_baton=True), + ] + + with patch("aare.daq.server.cfg", mock_cfg): + print(mock_backend["cfg"], id(mock_backend["cfg"])) + import aare.daq.server as server + print(server.cfg, id(server.cfg)) + response = client.get("/admin/gui_sessions", headers={"Authorization": "Bearer fake-token"}) + + assert response.status_code == 200 + payload = response.json() + assert payload[0]["holds_baton"] is False + assert payload[1]["holds_baton"] is True \ No newline at end of file diff --git a/tests/unit/daq/test_mount.py b/tests/unit/daq/test_mount.py new file mode 100644 index 00000000..3564022e --- /dev/null +++ b/tests/unit/daq/test_mount.py @@ -0,0 +1,263 @@ +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from unittest.mock import MagicMock + +from aareDB import SampleEventType + +from aare.common.models import DewarAddress, SampleShortInfo, DAQOperation +from aare.common.tell_models import TellActivityEnum, TellPhaseEnum, TellStateModel +from aare.daq.config import ABR_POS_MOUNT, BeamlineStateEnum +from aare.daq.daq import AareDAQ +from aare.daq.operations.mounting.models import MountingResult +from aare.daq.operations.mounting.service import MountingService +from aare.daq.operations.screenshot.service import ScreenshotService + + +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_daq(previous_sample: SampleShortInfo | None) -> AareDAQ: + daq = object.__new__(AareDAQ) + daq._AareDAQ__cfg = SimpleNamespace(current_sample=previous_sample) + daq._AareDAQ__aare = SimpleNamespace(send_sample_event=MagicMock()) + daq._AareDAQ__set_state = MagicMock() + daq._handle_operation_error = MagicMock() + daq.save_screenshot_db = MagicMock() + daq.sync_current_sample_from_tell = MagicMock(return_value=previous_sample) + daq._create_mounting_service = MagicMock() + return daq + + +def test_create_mounting_service_builds_expected_context(mock_logger): + daq = object.__new__(AareDAQ) + daq._AareDAQ__cfg = SimpleNamespace() + daq._AareDAQ__devs = SimpleNamespace() + + service = daq._create_mounting_service() + + assert isinstance(service, MountingService) + assert service.ctx.cfg is daq._AareDAQ__cfg + assert service.ctx.devs is daq._AareDAQ__devs + assert service.ctx.mount_position == ABR_POS_MOUNT + + +def test_was_previous_sample_unmounted_since_prefers_tell_phase_confirmation(): + previous_sample = _make_sample(1, "old") + daq = _make_daq(previous_sample) + + started_at = datetime.now(timezone.utc) - timedelta(seconds=5) + daq._safe_tell_state = MagicMock( + return_value=TellStateModel( + activity=TellActivityEnum.MOUNTING, + operation="mount", + phase=TellPhaseEnum.PICKING_NEW_SAMPLE, + last_update_ts=datetime.now(timezone.utc).isoformat(), + last_event_class="Motion Sync", + last_event_value="Sample get on Puck", + ) + ) + daq._get_tell_events_from_redis = MagicMock(return_value=[]) + + assert daq._was_previous_sample_unmounted_since(started_at) is True + + +def test_was_previous_sample_unmounted_since_falls_back_to_redis_event_history(): + previous_sample = _make_sample(1, "old") + daq = _make_daq(previous_sample) + + started_at = datetime.now(timezone.utc) - timedelta(seconds=5) + daq._safe_tell_state = MagicMock(return_value=None) + daq._get_tell_events_from_redis = MagicMock( + return_value=[ + { + "timestamp": (started_at - timedelta(seconds=2)).isoformat(), + "class": "Motion Sync", + "event": "Sample put on Puck", + }, + { + "timestamp": (started_at + timedelta(seconds=1)).isoformat(), + "class": "Motion Sync", + "event": "Sample put on Puck", + }, + ] + ) + + assert daq._was_previous_sample_unmounted_since(started_at) is True + + +def test_execute_mount_and_prepare_success_uses_mounting_result_fields(): + previous_sample = _make_sample(1, "old") + target_sample = _make_sample(2, "new") + daq = _make_daq(previous_sample) + + service = MagicMock() + service.execute.return_value = MountingResult( + success=True, + mounted_sample=target_sample, + previous_sample=previous_sample, + did_unmount_previous=True, + ) + daq._create_mounting_service.return_value = service + + result = daq._execute_mount_and_prepare(target_sample) + + assert result is True + + send_calls = daq._AareDAQ__aare.send_sample_event.call_args_list + assert send_calls[0].args[0] == previous_sample + assert send_calls[0].args[1] == SampleEventType.UNMOUNTING + assert send_calls[1].args[0] == target_sample + assert send_calls[1].args[1] == SampleEventType.MOUNTING + assert send_calls[2].args[0] == previous_sample + assert send_calls[2].args[1] == SampleEventType.UNMOUNTED + assert send_calls[3].args[0] == target_sample + assert send_calls[3].args[1] == SampleEventType.MOUNTED + + daq.save_screenshot_db.assert_called_once_with(target_sample.db_id, f"{target_sample.db_id}_mounted") + assert daq._AareDAQ__set_state.call_args_list[0].args[0] == BeamlineStateEnum.RobotSampleExchange + assert daq._AareDAQ__set_state.call_args_list[-1].args[0] == BeamlineStateEnum.SampleAlignment + + +def test_execute_mount_and_prepare_marks_previous_sample_unmounted_when_mount_fails_after_auto_unmount(): + previous_sample = _make_sample(1, "old") + target_sample = _make_sample(2, "new") + daq = _make_daq(previous_sample) + + service = MagicMock() + service.execute.return_value = MountingResult( + success=False, + error=RuntimeError("mount failed after unmount"), + comment="mount failed after unmount", + previous_sample=previous_sample, + ) + daq._create_mounting_service.return_value = service + daq._was_previous_sample_unmounted_since = MagicMock(return_value=True) + + result = daq._execute_mount_and_prepare(target_sample) + + assert result is False + assert daq._AareDAQ__cfg.current_sample is None + + send_calls = daq._AareDAQ__aare.send_sample_event.call_args_list + assert send_calls[0].args[0] == previous_sample + assert send_calls[0].args[1] == SampleEventType.UNMOUNTING + + assert send_calls[1].args[0] == target_sample + assert send_calls[1].args[1] == SampleEventType.MOUNTING + + assert send_calls[2].args[0] == previous_sample + assert send_calls[2].args[1] == SampleEventType.UNMOUNTED + assert send_calls[2].kwargs["comment"] == "Auto-unmount succeeded before mount failed" + + daq._handle_operation_error.assert_called_once() + assert daq._handle_operation_error.call_args.kwargs["operation"] == DAQOperation.MOUNT + assert daq._handle_operation_error.call_args.kwargs["sample"] == target_sample + assert daq._handle_operation_error.call_args.kwargs["event_type"] == SampleEventType.MOUNTFAILED + + assert daq._AareDAQ__set_state.call_args_list[-1].args[0] == BeamlineStateEnum.SampleAlignment + + +def test_execute_mount_and_prepare_does_not_mark_previous_sample_unmounted_when_not_confirmed(): + previous_sample = _make_sample(1, "old") + target_sample = _make_sample(2, "new") + daq = _make_daq(previous_sample) + + service = MagicMock() + service.execute.return_value = MountingResult( + success=False, + error=RuntimeError("mount failed before unmount confirmation"), + comment="mount failed before unmount confirmation", + previous_sample=previous_sample, + ) + daq._create_mounting_service.return_value = service + daq._was_previous_sample_unmounted_since = MagicMock(return_value=False) + + result = daq._execute_mount_and_prepare(target_sample) + + assert result is False + assert daq._AareDAQ__cfg.current_sample == previous_sample + + send_calls = daq._AareDAQ__aare.send_sample_event.call_args_list + assert len(send_calls) == 2 + assert send_calls[0].args[0] == previous_sample + assert send_calls[0].args[1] == SampleEventType.UNMOUNTING + assert send_calls[1].args[0] == target_sample + assert send_calls[1].args[1] == SampleEventType.MOUNTING + + daq._handle_operation_error.assert_called_once() + assert daq._handle_operation_error.call_args.kwargs["event_type"] == SampleEventType.MOUNTFAILED + + +def test_execute_mount_and_prepare_unmount_success_uses_unmount_operation(): + previous_sample = _make_sample(1, "old") + daq = _make_daq(previous_sample) + + service = MagicMock() + service.execute.return_value = MountingResult( + success=True, + mounted_sample=None, + previous_sample=previous_sample, + did_unmount_previous=True, + ) + daq._create_mounting_service.return_value = service + + result = daq._execute_mount_and_prepare(None) + + assert result is True + + send_calls = daq._AareDAQ__aare.send_sample_event.call_args_list + assert len(send_calls) == 2 + assert send_calls[0].args[0] == previous_sample + assert send_calls[0].args[1] == SampleEventType.UNMOUNTING + assert send_calls[1].args[0] == previous_sample + assert send_calls[1].args[1] == SampleEventType.UNMOUNTED + + daq.save_screenshot_db.assert_not_called() + assert daq._AareDAQ__set_state.call_args_list[0].args[0] == BeamlineStateEnum.RobotSampleExchange + assert daq._AareDAQ__set_state.call_args_list[-1].args[0] == BeamlineStateEnum.SampleAlignment + + +def test_create_loop_centering_service_uses_shared_screenshot_service(): + daq = object.__new__(AareDAQ) + daq._AareDAQ__cfg = SimpleNamespace() + daq._AareDAQ__devs = SimpleNamespace() + daq._AareDAQ__mlbox = SimpleNamespace(predict_all_best=MagicMock()) + daq._screenshot_service = MagicMock(spec=ScreenshotService) + daq._append_smargon_trace = MagicMock() + + type(daq).sample = property(lambda self: None) + type(daq).sample_geometry = property(lambda self: SimpleNamespace()) + type(daq).status = property(lambda self: SimpleNamespace()) + + service = daq._create_loop_centering_service() + + assert service.ctx.services.screenshots is daq._screenshot_service + + +def test_create_raster_service_uses_shared_screenshot_service(): + daq = object.__new__(AareDAQ) + daq._AareDAQ__cfg = SimpleNamespace() + daq._AareDAQ__devs = SimpleNamespace() + daq._AareDAQ__mlbox = SimpleNamespace() + daq._AareDAQ__jfjoch = SimpleNamespace() + daq._AareDAQ__aare = SimpleNamespace() + daq._AareDAQ__set_state = MagicMock() + daq._screenshot_service = MagicMock(spec=ScreenshotService) + + type(daq).sample = property(lambda self: None) + type(daq).sample_geometry = property(lambda self: SimpleNamespace()) + type(daq).status = property(lambda self: SimpleNamespace()) + + service = daq._create_raster_service() + + assert service.ctx.services.screenshots is daq._screenshot_service \ No newline at end of file diff --git a/tests/unit/gui/test_axis_video_thread.py b/tests/unit/gui/test_axis_video_thread.py new file mode 100644 index 00000000..2f3fd9d9 --- /dev/null +++ b/tests/unit/gui/test_axis_video_thread.py @@ -0,0 +1,70 @@ +import pytest +import numpy as np +import cv2 +from unittest.mock import MagicMock, patch +from PySide6.QtGui import QImage +from aare.gui.threads.axis_video_thread import VideoThread + +@pytest.fixture +def video_thread(qtbot): + thread = VideoThread("127.0.0.1", camera=1) + return thread + +def test_init(video_thread): + assert video_thread.camera_ip == "127.0.0.1" + assert video_thread.camera == 1 + assert video_thread.running is False + assert video_thread.is_busy is False + +def test_set_camera_ip(video_thread): + video_thread.set_camera_ip("192.168.1.1") + assert video_thread.camera_ip == "192.168.1.1" + +def test_set_busy(video_thread): + video_thread.set_busy(True) + assert video_thread.is_busy is True + +def test_process_buffer_success(video_thread, qtbot): + # Create a small valid JPEG + img = np.zeros((10, 10, 3), dtype=np.uint8) + _, jpeg_bytes = cv2.imencode('.jpg', img) + jpeg_bytes = jpeg_bytes.tobytes() + + boundary = b'--boundary' + buffer = b'--boundary\r\nContent-Type: image/jpeg\r\n\r\n' + jpeg_bytes + b'\r\n--boundary' + + with qtbot.waitSignal(video_thread.frame_ready, timeout=1000) as blocker: + video_thread._process_buffer(buffer, boundary) + + assert isinstance(blocker.args[0], QImage) + assert blocker.args[0].width() == 10 + assert blocker.args[0].height() == 10 + +def test_process_buffer_invalid_jpeg(video_thread): + boundary = b'--boundary' + buffer = b'--boundary\r\nContent-Type: image/jpeg\r\n\r\nNOT_A_JPEG\r\n--boundary' + + # Should not emit anything or raise exception + video_thread._process_buffer(buffer, boundary) + +def test_stop(video_thread): + video_thread.session = MagicMock() + # Mocking wait and quit to avoid actual thread blocking in test + video_thread.quit = MagicMock() + video_thread.wait = MagicMock() + + video_thread.stop() + assert video_thread.running is False + video_thread.session.close.assert_called_once() + video_thread.quit.assert_called_once() + video_thread.wait.assert_called_once() + +@patch('requests.Session') +def test_run_connection_error(mock_session_class, video_thread, qtbot): + mock_session = mock_session_class.return_value + mock_session.get.side_effect = Exception("Connection Refused") + + with qtbot.waitSignal(video_thread.error_occurred, timeout=1000) as blocker: + video_thread.run() + + assert "Unexpected error" in blocker.args[0] diff --git a/tests/unit/gui/test_camera_thread.py b/tests/unit/gui/test_camera_thread.py new file mode 100644 index 00000000..abe8d500 --- /dev/null +++ b/tests/unit/gui/test_camera_thread.py @@ -0,0 +1,105 @@ +import pytest +import json +import numpy as np +import cv2 +import zmq +from unittest.mock import MagicMock, patch +from PySide6.QtGui import QPixmap +from aare.gui.threads.camera_thread import SampleCameraThread +from aare.common.models import DAQStatusModel + +@pytest.fixture +def mock_zmq(): + with patch('zmq.Context') as mock_ctx_class: + mock_ctx = mock_ctx_class.return_value + mock_socket = mock_ctx.socket.return_value + yield mock_socket + +@pytest.fixture +def camera_thread(mock_zmq, qtbot): + thread = SampleCameraThread("tcp://127.0.0.1:5555") + return thread + +def test_init(camera_thread): + assert camera_thread.running is True + assert camera_thread._SampleCameraThread__camera_available is False + +def test_update_daq_status(camera_thread): + s = MagicMock() + s.geom.beam_location_pxl.x = 100 + s.geom.beam_location_pxl.y = 200 + camera_thread.update_daq_status(s) + assert camera_thread._SampleCameraThread__beam_x == 100 + assert camera_thread._SampleCameraThread__beam_y == 200 + +def test_enable_focus_measurement(camera_thread): + camera_thread.enable_focus_measurement(True) + assert camera_thread._SampleCameraThread__measure_focus is True + camera_thread.enable_focus_measurement(False) + assert camera_thread._SampleCameraThread__measure_focus is False + +def test_run_success(camera_thread, mock_zmq, qtbot): + # Create a small valid JPEG + img = np.zeros((10, 10, 3), dtype=np.uint8) + _, jpeg_bytes = cv2.imencode('.jpg', img) + + header = json.dumps({"encoding": "jpeg"}).encode("utf-8") + + def side_effect(): + camera_thread.running = False + return [header, jpeg_bytes.tobytes()] + + mock_zmq.recv_multipart.side_effect = side_effect + + # Increase timeout to 5000ms + with qtbot.waitSignal(camera_thread.camera_image, timeout=5000): + camera_thread.run() + + assert camera_thread._SampleCameraThread__camera_available is True + +def test_run_bayer_success(camera_thread, mock_zmq, qtbot): + # Create a small bayer image + h, w = 10, 10 + raw = np.zeros((h, w), dtype=np.uint8) + header = json.dumps({"shape": [h, w]}).encode("utf-8") + + def side_effect(): + camera_thread.running = False + return [header, raw.tobytes()] + + mock_zmq.recv_multipart.side_effect = side_effect + + # Increase timeout to 5000ms + with qtbot.waitSignal(camera_thread.camera_image, timeout=5000): + camera_thread.run() + + assert camera_thread._SampleCameraThread__camera_available is True + +def test_run_zmq_timeout(camera_thread, mock_zmq, qtbot): + def side_effect(): + camera_thread.running = False + raise zmq.Again() + + mock_zmq.recv_multipart.side_effect = side_effect + + # We want to check if it emits nan or 0 for FPS on timeout + # But it only emits every 0.5s. + camera_thread._SampleCameraThread__fps_emit_period_s = 0.0 + + # Increase timeout to 5000ms + with qtbot.waitSignal(camera_thread.fps_measure, timeout=5000): + try: + camera_thread.run() + except StopIteration: + pass + +def test_stop(camera_thread, mock_zmq): + # Mocking wait and quit to avoid actual thread blocking in test + camera_thread.quit = MagicMock() + camera_thread.wait = MagicMock() + + camera_thread.stop() + assert camera_thread.running is False + mock_zmq.close.assert_called_once() + camera_thread.quit.assert_called_once() + camera_thread.wait.assert_called_once() diff --git a/tests/unit/gui/test_message_box.py b/tests/unit/gui/test_message_box.py new file mode 100644 index 00000000..bf8ca4ea --- /dev/null +++ b/tests/unit/gui/test_message_box.py @@ -0,0 +1,62 @@ +import pytest +from unittest.mock import MagicMock, patch +from PySide6.QtWidgets import QMessageBox +from aare.gui.widgets.message_box import reply_box, timer_box, ring_current_low_check, experiment_hutch_shutter_check, ring_current_auto_check + +def test_reply_box(qtbot): + parent = MagicMock() + with patch('PySide6.QtWidgets.QMessageBox.question') as mock_question: + mock_question.return_value = QMessageBox.StandardButton.Yes + res = reply_box(parent, "Title", "Message") + assert res == QMessageBox.StandardButton.Yes + mock_question.assert_called_once() + +def test_timer_box_auto_accept(qtbot): + parent = None + condition_func = MagicMock(return_value=True) + + # We need to process events for timer to fire + box = timer_box(parent, condition_func=condition_func) + + # Wait until box is closed by check() + qtbot.waitUntil(lambda: not box.isVisible(), timeout=2000) + + assert box.result() == QMessageBox.StandardButton.Yes + +def test_ring_current_low_check_ok(qtbot): + # Should return True immediately if current is high enough + assert ring_current_low_check(None, 200.0) is True + +def test_ring_current_low_check_low_yes(qtbot): + with patch('aare.gui.widgets.message_box.reply_box') as mock_reply: + mock_reply.return_value = QMessageBox.StandardButton.Yes + assert ring_current_low_check(None, 50.0) is True + mock_reply.assert_called_once() + +def test_ring_current_low_check_low_no(qtbot): + with patch('aare.gui.widgets.message_box.reply_box') as mock_reply: + mock_reply.return_value = QMessageBox.StandardButton.No + assert ring_current_low_check(None, 50.0) is False + +def test_experiment_hutch_shutter_check_open(qtbot): + assert experiment_hutch_shutter_check(None, True) is True + +def test_experiment_hutch_shutter_check_closed_yes(qtbot): + with patch('aare.gui.widgets.message_box.reply_box') as mock_reply: + mock_reply.return_value = QMessageBox.StandardButton.Yes + assert experiment_hutch_shutter_check(None, False) is True + +def test_ring_current_auto_check_yes(qtbot): + # This one uses a nested event loop, which can be tricky to test. + # We'll mock timer_box to return a box that we can close manually. + with patch('aare.gui.widgets.message_box.timer_box') as mock_timer_box: + box = QMessageBox() + box.setStandardButtons(QMessageBox.StandardButton.Yes) + mock_timer_box.return_value = box + + # We need to close the box after some time to break the loop + from PySide6.QtCore import QTimer + QTimer.singleShot(100, lambda: box.done(QMessageBox.StandardButton.Yes)) + + res = ring_current_auto_check(None, 50.0, lambda: False) + assert res is True