From 3af118fb644f52ac72d9bf6f5ba815cbf75c6a5e Mon Sep 17 00:00:00 2001
From: perl_d <david.perl@psi.ch>
Date: Fri, 7 Aug 2026 15:12:12 +0200
Subject: [PATCH] wrap mlbox

---
 .../beamline_dispatch/default/beamline_dispatch.py  |  5 +++++
 src/aare/beamline_dispatch/protocols.py             |  4 ++++
 .../simulated/beamline_dispatch.py                  | 13 +++++++++++--
 src/aare/daq/daq.py                                 |  4 +++-
 4 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/src/aare/beamline_dispatch/default/beamline_dispatch.py b/src/aare/beamline_dispatch/default/beamline_dispatch.py
index a2830190..2caf493d 100644
--- a/src/aare/beamline_dispatch/default/beamline_dispatch.py
+++ b/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)
diff --git a/src/aare/beamline_dispatch/protocols.py b/src/aare/beamline_dispatch/protocols.py
index 61f2ebcc..fd649bcc 100644
--- a/src/aare/beamline_dispatch/protocols.py
+++ b/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: ...
diff --git a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py
index 6ff05f59..01634238 100644
--- a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py
+++ b/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()
diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py
index bf6acedc..27b5f94b 100644
--- a/src/aare/daq/daq.py
+++ b/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