diff --git a/src/aare/daq/mlbox.py b/src/aare/daq/mlbox.py index f70b67fb..26417070 100644 --- a/src/aare/daq/mlbox.py +++ b/src/aare/daq/mlbox.py @@ -1,14 +1,15 @@ from dataclasses import dataclass, field - -from typing import Optional, Iterable +from typing import Optional, Iterable, TYPE_CHECKING, Any import cv2 import numpy as np import time -from aarelcinfer_client.models import LatestPredictionModel +if TYPE_CHECKING: + from aarelcinfer_client.models import LatestPredictionModel +else: + LatestPredictionModel = Any -from aare.common.aarelc_infer import AareLCInferWrapper from aare.common.beamline import MXBeamline from aare.common.models import MLBoxModel, MLOutputModel, MLBoxType, BoundingBoxModel from aare.common.logger_config import setup_logger @@ -47,23 +48,37 @@ class MLBundle: class MlBox: - RETRY_COUNT = 3 RETRY_SLEEP_S = 0.1 + def __init__(self, bl: MXBeamline, wrapper=None): + """ + wrapper: + Optional injected inference wrapper for tests. + If None, the production wrapper is created lazily. + """ + self.__beamline = bl + + if wrapper is not None: + self.__wrapper = wrapper + return - def __init__(self, bl:MXBeamline): #mx-aare-test.psi.ch, mx-ml.psi.ch if bl == MXBeamline.SIMULATED: raise NotImplementedError(f"MLBox bundle mode not implemented for {bl}") + elif bl == MXBeamline.X06DA: raise NotImplementedError(f"MLBox bundle mode not implemented for {bl}") + elif bl == MXBeamline.X10SA: + # Lazy import keeps module importable without external client installed + from aare.common.aarelc_infer import AareLCInferWrapper self.__wrapper = AareLCInferWrapper(bl) - self.__beamline = bl + elif bl == MXBeamline.X06SA: raise NotImplementedError(f"MLBox bundle mode not implemented for {bl}") + else: - raise Exception(f"unknown beamline {bl}") + raise ValueError(f"unknown beamline {bl}") @staticmethod def _decode_bundle_image(jpeg_bytes: bytes | None) -> np.ndarray | None: diff --git a/src/aare/daq/workflows.py b/src/aare/daq/workflows.py index 50630e1e..bda7fcae 100644 --- a/src/aare/daq/workflows.py +++ b/src/aare/daq/workflows.py @@ -1,10 +1,10 @@ from aare.common.models import StagePositionEnum from aare.daq.devices import BeamlineDevices from aare.daq.config import BeamlineConfig, ABR_POS_MOUNT -from pxii_bec.macros.planner import BeamlineState from aare.devices.area_detector import AutoEnum from aare.common.logger_config import setup_logger +from aare.devices.bec_worker import BeamlineState logger = setup_logger("aareDAQ") diff --git a/src/aare/devices/bec_worker.py b/src/aare/devices/bec_worker.py index 041cfa29..c30cc3ae 100644 --- a/src/aare/devices/bec_worker.py +++ b/src/aare/devices/bec_worker.py @@ -1,8 +1,9 @@ +from enum import Enum + from bec_ipython_client import BECIPythonClient from bec_ipython_client.signals import OperationMode from bec_lib.service_config import ServiceConfig from bec_lib.procedures.helper import FrontendProcedureHelper, BackendProcedureHelper -from pxii_bec.macros.planner import BeamlineState from aare.common.beamline import MXBeamline, mx_beamline @@ -23,6 +24,20 @@ from aare.common.beamline import MXBeamline, mx_beamline def bec_exception_handler(exception: Exception): print(f"Exception: {exception}") + +class BeamlineState(str, Enum): + ROBOT_SAMPLE_EXCHANGE = "robot_sample_exchange" + SAMPLE_ALIGNMENT = "sample_alignment" + DATA_COLLECTION = "data_collection" + DC_XRF = "DC_XRF" + MANUAL_SAMPLE_EXCHANGE = "manual_sample_exchange" + BEAM_VISUALISATION = "beam_visualisation" + FLUX_MEASUREMENT = "flux_measurement" + BEAMSTOP_ALIGNMENT = "beamstop_alignment" + MAINTENANCE = "maintenance" + XTAL_SNAPSHOT = "xtal_snapshot" + + class BECClientWorker: def __init__(self, beamline:MXBeamline, name:str = "default"): BEAMLINE = beamline.value.lower() diff --git a/tests/conftest.py b/tests/conftest.py index a9c175a3..571618ac 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,5 @@ import sys +from unittest.mock import MagicMock import pytest @@ -6,6 +7,30 @@ from PySide6.QtWidgets import QApplication from aare.common.models import SampleShortInfo, DewarAddress +# external service mocks +from unittest.mock import MagicMock +import sys + +def mock_module(name): + m = MagicMock() + m.__path__ = [] + sys.modules[name] = m + return m + +mock_module("aarelcinfer_client") +mock_module("aarelcinfer_client.api") +mock_module("aarelcinfer_client.models") + +mock_module("aarescan_client") +mock_module("aarescan_client.models") +mock_module("aarescan_client.models.grid_request") +mock_module("aarescan_client.models.screen_request") + +sys.modules.setdefault("pxii_bec", MagicMock()) +sys.modules.setdefault("pxii_bec.macros", MagicMock()) +sys.modules.setdefault("pxii_bec.macros.planner", MagicMock()) + + @pytest.fixture def sample_info(): diff --git a/tests/unit/daq/test_mlbox_logic.py b/tests/unit/daq/test_mlbox_logic.py deleted file mode 100644 index 6f2c23c7..00000000 --- a/tests/unit/daq/test_mlbox_logic.py +++ /dev/null @@ -1,40 +0,0 @@ -from aare.daq.mlbox import MlBox -from aare.common.models import MLOutputModel, MLBoxType - - -def test_preferred_class_prefers_loop_over_pin_when_margin_not_exceeded(): - boxes = MLOutputModel() - boxes.add_box(MLBoxType.Pin, (0, 0, 10, 10), 0.8) - boxes.add_box(MLBoxType.Loop_face, (1, 1, 9, 9), 0.75) - - best = MlBox.get_preferred_class_box_with_confidence_threshold( - boxes, - loop_preference_margin=0.1, - ) - - assert best is not None - assert best.cls == MLBoxType.Loop_face - -def test_prediction_score_empty(): - assert MlBox._prediction_score(None) == (0, 0.0) - - -def test_prediction_score_uses_count_and_max_confidence(): - model = MLOutputModel() - model.add_box(MLBoxType.Pin, (0, 0, 10, 10), 0.4) - model.add_box(MLBoxType.Crystal, (1, 1, 9, 9), 0.8) - - assert MlBox._prediction_score(model) == (2, 0.8) - -def test_best_by_class_keeps_highest_confidence_per_class(): - results = [ - {"class": 1, "confidence": 0.4, "box": {"x1": 0, "y1": 0, "x2": 10, "y2": 10}}, - {"class": 1, "confidence": 0.8, "box": {"x1": 1, "y1": 1, "x2": 11, "y2": 11}}, - {"class": 2, "confidence": 0.6, "box": {"x1": 2, "y1": 2, "x2": 12, "y2": 12}}, - ] - - out = MlBox._best_by_class(results) - - assert out is not None - assert out.get_best_for_class(MLBoxType.Pin).conf == 0.8 - assert out.get_best_for_class(MLBoxType.Crystal).conf == 0.6 diff --git a/tests/unit/daq/test_server.py b/tests/unit/daq/test_server.py index c665988c..071dfdac 100644 --- a/tests/unit/daq/test_server.py +++ b/tests/unit/daq/test_server.py @@ -10,6 +10,7 @@ import numpy as np os.environ["JWT_AAREDAQ_KEY"] = "test_key_for_unit_testing" # Mock the entire backend before importing the app to prevent any real initialisation +import aare.daq.daq with patch("aare.daq.daq.AareDAQ"), \ patch("aare.daq.config.MXBeamline"), \ patch("aare.daq.config.BeamlineConfig"), \