tests: more tests

This commit is contained in:
2026-05-22 18:05:46 +02:00
parent 0097b246bf
commit 3654e126e6
8 changed files with 1136 additions and 0 deletions
+199
View File
@@ -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"
@@ -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}]
@@ -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")
+203
View File
@@ -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
+263
View File
@@ -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
+70
View File
@@ -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]
+105
View File
@@ -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()
+62
View File
@@ -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