wrap mlbox

src/aare/beamline_dispatch/default/beamline_dispatch.py

@@ -8,6 +8,7 @@ from bec_ipython_client.signals import OperationMode
 from bec_lib.client import ServiceConfig
 
 from aare.beamline_dispatch.protocols import AuthDispatch, BeamlineDispatch, BecSetup, ConfigKeys
+from aare.daq.mlbox import MlBox
 
 logger = setup_logger("aareDAQ")
 
@@ -71,6 +72,7 @@ class DefaultDispatch(BeamlineDispatch, ABC):
     Should be safe and fail rather than assuming anything."""
 
     def __init__(self, beamline: MXBeamline) -> None:
+        self._beamline = beamline
         self._auth = DefaultAuthDispatch()
         self._config = DefaultConfigKeys(beamline)
         self._bec_setup = DefaultBecSetup(beamline)
@@ -86,3 +88,6 @@ class DefaultDispatch(BeamlineDispatch, ABC):
     @property
     def bec_setup(self):
         return self._bec_setup
+
+    def get_mlbox(self) -> MlBox:
+        return MlBox(self._beamline)

src/aare/beamline_dispatch/protocols.py

@@ -4,6 +4,8 @@ from typing import Any
 from aarecommon.models.models import MXBeamline
 from bec_ipython_client import BECIPythonClient
 
+from aare.daq.mlbox import MlBox
+
 
 class AuthDispatch(ABC):
     @abstractmethod
@@ -63,3 +65,5 @@ class BeamlineDispatch(ABC):
     @property
     @abstractmethod
     def config(self) -> ConfigKeys: ...
+    @abstractmethod
+    def get_mlbox(self) -> MlBox: ...

src/aare/beamline_dispatch/simulated/beamline_dispatch.py

@@ -1,14 +1,20 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 from unittest.mock import MagicMock
 
 from aarecommon.config.beamline import MXBeamline
-from bec_ipython_client import BECIPythonClient
 from bec_lib.devicemanager import DeviceContainer
 
 from aare.beamline_dispatch.default.beamline_dispatch import DefaultConfigKeys, DefaultDispatch
 from aare.beamline_dispatch.protocols import AuthDispatch, BecMacros, BecSetup
 from aare.devices.bec_definitions import BecDevices
 
+if TYPE_CHECKING:
+    from bec_ipython_client import BECIPythonClient
+
+    from aare.daq.mlbox import MlBox
+
 
 class SimulatedAuthDispatch(AuthDispatch):
     def get_jwt_key(self) -> str:
@@ -57,3 +63,6 @@ class SimulatedDispatch(DefaultDispatch):
     @property
     def bec_macros(self):
         return self._bec_macros
+
+    def get_mlbox(self) -> MlBox:
+        return MagicMock()

src/aare/daq/daq.py

@@ -67,6 +67,7 @@ from aarecommon.models.rotation_scan import CompletedRotationScan, RotationScanR
 from aarecommon.models.tell import TellPhaseEnum, TellStateModel
 from aareDB import SampleEventType
 
+from aare.beamline_dispatch.beamline_dispatch import get_beamline_dispatch
 from aare.daq import workflows
 from aare.daq.aaredb import AareWrapper
 from aare.daq.config import ABR_POS_MOUNT, BeamlineConfig, BeamlineStateEnum
@@ -273,8 +274,9 @@ class AareDAQ:
     def __init__(self, cfg: BeamlineConfig, bl: MXBeamline):
         self.last_time = 0.0
         self._cfg = cfg
+        self._dispatch = get_beamline_dispatch()
+        self._mlbox = self._dispatch.get_mlbox()
         self._devs = BeamlineDevices(bl)
-        self._mlbox = MlBox(bl)
         self._jfjoch = JFJochWrapper(bl)
         self._bl = bl.value.upper()
         self._beamline = bl