bec-server 4.1.0__tar.gz → 4.1.2__tar.gz

{bec_server-4.1.0 → bec_server-4.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: bec-server
-Version: 4.1.0
+Version: 4.1.2
 Summary: BEC server
 Project-URL: Bug Tracker, https://github.com/bec-project/bec/issues
 Project-URL: Homepage, https://github.com/bec-project/bec

{bec_server-4.1.0 → bec_server-4.1.2}/bec_server/file_writer/file_writer_manager.py

@@ -4,7 +4,7 @@ import os
 import threading
 import time
 import traceback
-from collections import defaultdict
+from collections import defaultdict, deque
 
 from bec_lib import messages, plugin_helper
 from bec_lib.alarm_handler import Alarms
@@ -60,12 +60,23 @@ class ScanStorage:
         """
         self.scan_segments[point_id] = data
 
+    @property
+    def baseline_ready(self) -> bool:
+        """Whether no baseline read was requested or its bundled data has arrived."""
+        if self.status_msg is None:
+            return True
+        return not self.status_msg.info.get("baseline_readout_requested", False) or bool(
+            self.baseline
+        )
+
     def ready_to_write(self) -> bool:
         """
         Check if the scan is ready to be written to file.
         """
         if self.forced_finish:
             return True
+        if not self.baseline_ready:
+            return False
         if self.enforce_sync:
             # wait for all points to be received. Since this method will be called for every
             # update of the scan segments, we can also accept to write after the scan is finished
@@ -106,8 +117,16 @@ class FileWriterManager(BECService):
         self.available_beamline_states: list[messages.BeamlineStateConfig] = []
         self.beamline_state_subscriptions: set[str] = set()
         self.beamline_states: dict[str, messages.BeamlineStateMessage] = {}
+        self.scan_storage: dict[str, ScanStorage] = {}
+        self._pending_baseline_writes: dict[str, dict] = {}
+        self._pending_file_references: dict[str, dict[str, messages.FileMessage]] = {}
+        self._finished_scan_ids: deque[str] = deque(maxlen=1000)
         self.connector.register(MessageEndpoints.scan_segment(), cb=self._scan_segment_callback)
         self.connector.register(MessageEndpoints.scan_status(), cb=self._scan_status_callback)
+        self.connector.register(MessageEndpoints.scan_baseline(), cb=self._scan_baseline_callback)
+        self.connector.register(
+            patterns=MessageEndpoints.public_file("*", "*"), cb=self._file_reference_callback
+        )
         self.connector.register(
             patterns=MessageEndpoints.device_read_configuration("*"),  # type: ignore
             cb=self._device_configuration_callback,
@@ -122,7 +141,6 @@ class FileWriterManager(BECService):
         )
         self._run_storage_copy_plugin = self._load_storage_copy_plugin()
         self.async_writer = None
-        self.scan_storage: dict[str, ScanStorage] = {}
         self.file_writer = HDF5FileWriter(self)
         self.status = messages.BECStatus.RUNNING
         self.refresh_device_configs()
@@ -141,6 +159,43 @@ class FileWriterManager(BECService):
         msg = msg.value
         self.update_scan_storage_with_status(msg)
 
+    def _scan_baseline_callback(self, msg: MessageObject) -> None:
+        baseline_msg = msg.value
+        with self._lock:
+            if baseline_msg.scan_id in self._finished_scan_ids:
+                return
+            storage = self.scan_storage.get(baseline_msg.scan_id)
+            if storage is None:
+                self._pending_baseline_writes[baseline_msg.scan_id] = baseline_msg.data
+                return
+            storage.baseline = baseline_msg.data
+            self.check_storage_status(scan_id=baseline_msg.scan_id)
+
+    def _file_reference_callback(self, msg: MessageObject) -> None:
+        _, scan_id, _, name = msg.topic.rsplit("/", 3)
+        file_msg = msg.value
+        if name == "master" or file_msg.is_master_file:
+            return
+        with self._lock:
+            if scan_id in self._finished_scan_ids:
+                return
+            storage = self.scan_storage.get(scan_id)
+            if storage is None:
+                self._pending_file_references.setdefault(scan_id, {})[name] = file_msg
+                return
+            storage.file_references[name] = file_msg
+
+    def _get_or_create_scan_storage(self, scan_id: str, scan_number: int | None) -> ScanStorage:
+        """Create scan storage and transfer any notifications received before its status or data."""
+        with self._lock:
+            storage = self.scan_storage.get(scan_id)
+            if storage is None:
+                storage = ScanStorage(scan_number=scan_number, scan_id=scan_id)
+                storage.baseline = self._pending_baseline_writes.pop(scan_id, {})
+                storage.file_references = self._pending_file_references.pop(scan_id, {})
+                self.scan_storage[scan_id] = storage
+            return storage
+
     def _device_configuration_callback(self, msg):
         topic, msg = msg.topic, msg.value
         device = topic.split("/")[-1]
@@ -255,10 +310,7 @@ class FileWriterManager(BECService):
         if scan_id is None:
             return
 
-        if not self.scan_storage.get(scan_id):
-            self.scan_storage[scan_id] = ScanStorage(
-                scan_number=msg.content["info"].get("scan_number"), scan_id=scan_id
-            )
+        self._get_or_create_scan_storage(scan_id, msg.content["info"].get("scan_number"))
 
         for state in self.beamline_states.values():
             if state.name not in self.scan_storage[scan_id].beamline_states:
@@ -324,60 +376,13 @@ class FileWriterManager(BECService):
         scan_id = msg.content.get("scan_id")
         if scan_id is None:
             return
-        if not self.scan_storage.get(scan_id):
-            self.scan_storage[scan_id] = ScanStorage(
-                scan_number=msg.metadata.get("scan_number"), scan_id=scan_id
-            )
+        self._get_or_create_scan_storage(scan_id, msg.metadata.get("scan_number"))
         self.scan_storage[scan_id].append(
             point_id=msg.content.get("point_id"), data=msg.content.get("data")
         )
         logger.debug(msg.content.get("point_id"))
         self.check_storage_status(scan_id=scan_id)
 
-    def update_baseline_reading(self, scan_id: str) -> None:
-        """
-        Update the baseline reading for the scan.
-
-        Args:
-            scan_id (str): Scan ID
-        """
-        if not self.scan_storage.get(scan_id):
-            return
-        if self.scan_storage[scan_id].baseline:
-            return
-        baseline = self.connector.get(MessageEndpoints.public_scan_baseline(scan_id))
-        if not baseline:
-            return
-        self.scan_storage[scan_id].baseline = baseline.content["data"]
-        return
-
-    def update_file_references(self, scan_id: str) -> None:
-        """
-        Update the file references for the scan.
-        All external files ought to be announced to the endpoint public_file before the scan finishes. This function
-        retrieves the file references and adds them to the scan storage.
-
-        Args:
-            scan_id (str): Scan ID
-        """
-        if not self.scan_storage.get(scan_id):
-            return
-        msgs = self.connector.keys(MessageEndpoints.public_file(scan_id, "*"))
-        if not msgs:
-            return
-
-        # extract name from 'public/<scan_id>/file/<name>'
-        names = [msg.decode().split("/")[-1] for msg in msgs]
-        file_msgs = [
-            self.connector.get(MessageEndpoints.public_file(scan_id=scan_id, name=name))
-            for name in names
-        ]
-        if not file_msgs:
-            return
-        for name, file_msg in zip(names, file_msgs):
-            self.scan_storage[scan_id].file_references[name] = file_msg
-        return
-
     def update_device_configuration(self, device: str, msg: messages.DeviceMessage) -> None:
         """
         Update the device configuration. Note that this is a global configuration and not specific to a scan.
@@ -397,8 +402,6 @@ class FileWriterManager(BECService):
         with self._lock:
             if not self.scan_storage.get(scan_id):
                 return
-            self.update_baseline_reading(scan_id)
-            self.update_file_references(scan_id)
             if self.scan_storage[scan_id].ready_to_write():
                 self.write_file(scan_id)
 
@@ -481,6 +484,7 @@ class FileWriterManager(BECService):
         logger.info(f"Writing to file {file_path} took {time.time() - start_time:.2f} seconds.")
 
         self.scan_storage.pop(scan_id)
+        self._finished_scan_ids.append(scan_id)
         self.connector.set_and_publish(
             MessageEndpoints.public_file(scan_id, "master"),
             messages.FileMessage(file_path=file_path, done=True, successful=successful),

{bec_server-4.1.0 → bec_server-4.1.2}/bec_server/scan_server/scans/scan_actions.py

@@ -614,6 +614,7 @@ class ScanActions:
             metadata={"device_instr_id": status._device_instr_id, "readout_priority": "baseline"},
         )
         self._send(instr)
+        self._scan.scan_info.baseline_readout_requested = True
         if wait:
             status.wait()
         return status

{bec_server-4.1.0 → bec_server-4.1.2}/bec_server/scan_server/scans/scan_base.py

@@ -116,6 +116,9 @@ class ScanInfo(BaseModel):
     metadata: Annotated[dict, Field(description="Additional metadata for the scan")] = {}
 
     # progress tracking
+    baseline_readout_requested: Annotated[
+        bool, Field(description="Whether a read of nonempty baseline devices was requested.")
+    ] = False
     num_monitored_readouts: Annotated[
         int,
         Field(

{bec_server-4.1.0 → bec_server-4.1.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "bec-server"
-version = "4.1.0"
+version = "4.1.2"
 description = "BEC server"
 requires-python = ">=3.11"
 classifiers = [

{bec_server-4.1.0 → bec_server-4.1.2}/tests/tests_file_writer/test_file_writer_manager.py

@@ -1,5 +1,6 @@
 # pylint: skip-file
 import os
+import threading
 import time
 from unittest import mock
 
@@ -201,15 +202,11 @@ def test_check_storage_status(file_writer_manager_mock, scan_storage_mock):
 
     with (
         mock.patch.object(scan_storage_mock, "ready_to_write") as mock_ready_to_write,
-        mock.patch.object(file_manager, "update_baseline_reading") as mock_update_baseline,
-        mock.patch.object(file_manager, "update_file_references") as mock_update_file_references,
         mock.patch.object(file_manager, "write_file") as mock_write_file,
     ):
         mock_ready_to_write.return_value = True
         file_manager.check_storage_status(scan_id="scan_id")
         assert mock_ready_to_write.called
-        assert mock_update_baseline.call_args == mock.call("scan_id")
-        assert mock_update_file_references.call_args == mock.call("scan_id")
         assert mock_write_file.call_args == mock.call("scan_id")
 
 
@@ -331,18 +328,6 @@ def test_write_file_renames_tmp_file_on_exception(file_writer_manager_mock, scan
             mock_rename.assert_called_once_with(tmp_file_path, final_file_path)
 
 
-def test_update_baseline_reading(file_writer_manager_mock, scan_storage_mock):
-    file_manager = file_writer_manager_mock
-    file_manager.scan_storage["scan_id"] = scan_storage_mock
-    with mock.patch.object(file_manager, "connector") as mock_connector:
-        mock_connector.get.return_value = messages.ScanBaselineMessage(
-            scan_id="scan_id", data={"data": "data"}
-        )
-        file_manager.update_baseline_reading("scan_id")
-        assert file_manager.scan_storage["scan_id"].baseline == {"data": "data"}
-        mock_connector.get.assert_called_once_with(MessageEndpoints.public_scan_baseline("scan_id"))
-
-
 def test_scan_storage_append(scan_storage_mock):
     storage = scan_storage_mock
     storage.append(1, {"data": "data"})
@@ -358,21 +343,6 @@ def test_scan_storage_ready_to_write(scan_storage_mock):
     assert storage.ready_to_write() is True
 
 
-def test_update_file_references(file_writer_manager_mock):
-    file_manager = file_writer_manager_mock
-    with mock.patch.object(file_manager, "connector") as mock_connector:
-        file_manager.update_file_references("scan_id")
-        mock_connector.keys.assert_not_called()
-
-
-def test_update_file_references_gets_keys(file_writer_manager_mock, scan_storage_mock):
-    file_manager = file_writer_manager_mock
-    file_manager.scan_storage["scan_id"] = scan_storage_mock
-    with mock.patch.object(file_manager, "connector") as mock_connector:
-        file_manager.update_file_references("scan_id")
-        mock_connector.keys.assert_called_once_with(MessageEndpoints.public_file("scan_id", "*"))
-
-
 def test_update_scan_storage_with_status_ignores_none(file_writer_manager_mock):
     file_manager = file_writer_manager_mock
     file_manager.update_scan_storage_with_status(
@@ -397,11 +367,7 @@ def test_update_scan_storage_with_status_waits_for_v4_sync_segments(file_writer_
         info={"scan_number": 1, "monitor_sync": None},
     )
 
-    with (
-        mock.patch.object(file_manager, "update_baseline_reading"),
-        mock.patch.object(file_manager, "update_file_references"),
-        mock.patch.object(file_manager, "write_file") as write_file,
-    ):
+    with (mock.patch.object(file_manager, "write_file") as write_file,):
         file_manager.update_scan_storage_with_status(msg)
 
     assert storage.enforce_sync is True
@@ -559,3 +525,317 @@ def test_file_writer_manager_removes_beamline_state_subscription(file_writer_man
     file_manager._update_available_beamline_states({"data": msg})
     assert "State1" not in file_manager.beamline_state_subscriptions
     assert "State1" not in file_manager.beamline_states
+
+
+@pytest.mark.parametrize("enforce_sync,monitored", [(True, ["samx"]), (True, []), (False, [])])
+def test_ready_to_write_waits_for_baseline(enforce_sync, monitored):
+    storage = ScanStorage(1, "scan_id")
+    storage.status_msg = messages.ScanStatusMessage(
+        scan_id="scan_id",
+        status="closed",
+        info={"baseline_readout_requested": True},
+        readout_priority={"monitored": monitored, "baseline": ["samz"]},
+    )
+    storage.enforce_sync = enforce_sync
+    storage.scan_finished = True
+    storage.num_monitored_readouts = 0
+
+    assert storage.baseline_ready is False
+    assert storage.ready_to_write() is False
+
+    storage.baseline = {"samz": {"samz": {"value": 42, "timestamp": 1}}}
+
+    assert storage.baseline_ready is True
+    assert storage.ready_to_write() is True
+
+
+@pytest.mark.parametrize("arrival_order", ["baseline_first", "baseline_last"])
+def test_baseline_arrival_finalizes_scan_once(file_writer_manager_mock, arrival_order):
+    manager = file_writer_manager_mock
+    manager.scan_storage["scan_id"] = ScanStorage(1, "scan_id")
+    baseline_msg = messages.ScanBaselineMessage(
+        scan_id="scan_id", data={"samz": {"samz": {"value": 42, "timestamp": 1}}}
+    )
+    baseline_event = MessageObject(
+        topic=MessageEndpoints.scan_baseline().endpoint, value=baseline_msg
+    )
+    closed_msg = messages.ScanStatusMessage(
+        scan_id="scan_id",
+        status="closed",
+        scan_number=1,
+        scan_type="software_triggered",
+        num_points=1,
+        num_monitored_readouts=1,
+        readout_priority={"monitored": ["samx"], "baseline": ["samz"]},
+        info={"scan_number": 1, "monitor_sync": "bec", "baseline_readout_requested": True},
+    )
+    segment = messages.ScanMessage(
+        scan_id="scan_id", point_id=0, data={"samx": {"samx": {"value": 1, "timestamp": 1}}}
+    )
+    written_baselines = []
+
+    def write_file(scan_id):
+        written_baselines.append(manager.scan_storage.pop(scan_id).baseline)
+        manager._finished_scan_ids.append(scan_id)
+
+    with (
+        mock.patch.object(manager.connector, "get", return_value=None) as get,
+        mock.patch.object(manager, "write_file", side_effect=write_file) as write,
+    ):
+        if arrival_order == "baseline_first":
+            manager._scan_baseline_callback(baseline_event)
+            write.assert_not_called()
+
+        manager.insert_to_scan_storage(segment)
+        write.assert_not_called()
+        manager.update_scan_storage_with_status(closed_msg)
+
+        if arrival_order == "baseline_last":
+            write.assert_not_called()
+            manager._scan_baseline_callback(baseline_event)
+
+        write.assert_called_once_with("scan_id")
+        assert written_baselines == [baseline_msg.data]
+        get.assert_not_called()
+        assert "scan_id" not in manager.scan_storage
+
+        manager._scan_baseline_callback(baseline_event)
+        write.assert_called_once_with("scan_id")
+        assert "scan_id" not in manager.scan_storage
+        assert manager._pending_baseline_writes == {}
+
+
+def test_baseline_callback_does_not_create_scan_storage(file_writer_manager_mock):
+    manager = file_writer_manager_mock
+    baseline = messages.ScanBaselineMessage(scan_id="unknown_scan", data={"samz": {}})
+    with mock.patch.object(manager, "write_file") as write:
+        manager._scan_baseline_callback(
+            MessageObject(topic=MessageEndpoints.scan_baseline().endpoint, value=baseline)
+        )
+    write.assert_not_called()
+    assert manager.scan_storage == {}
+    assert manager._pending_baseline_writes == {"unknown_scan": baseline.data}
+
+
+@pytest.mark.parametrize("status", ["aborted", "halted", "user_completed"])
+def test_interrupted_scan_does_not_wait_for_baseline(file_writer_manager_mock, status):
+    manager = file_writer_manager_mock
+    storage = ScanStorage(1, "scan_id")
+    manager.scan_storage["scan_id"] = storage
+    status_msg = messages.ScanStatusMessage(
+        scan_id="scan_id",
+        status=status,
+        scan_number=1,
+        scan_type="software_triggered",
+        num_monitored_readouts=1,
+        readout_priority={"monitored": ["samx"], "baseline": ["samz"]},
+        info={"scan_number": 1, "monitor_sync": "bec", "baseline_readout_requested": True},
+    )
+    with (
+        mock.patch.object(manager.connector, "get", return_value=None),
+        mock.patch.object(manager, "write_file") as write,
+    ):
+        manager.update_scan_storage_with_status(status_msg)
+        assert storage.baseline_ready is False
+        write.assert_called_once_with("scan_id")
+
+
+@pytest.mark.parametrize("baseline_devices", [[], ["samz"]])
+@pytest.mark.parametrize("requested_info", [{}, {"baseline_readout_requested": False}])
+def test_scan_without_baseline_request_does_not_wait_for_baseline(
+    file_writer_manager_mock, baseline_devices, requested_info
+):
+    manager = file_writer_manager_mock
+    manager.scan_storage["scan_id"] = ScanStorage(1, "scan_id")
+    closed_msg = messages.ScanStatusMessage(
+        scan_id="scan_id",
+        status="closed",
+        scan_number=1,
+        scan_type="software_triggered",
+        num_monitored_readouts=0,
+        readout_priority={"monitored": [], "baseline": baseline_devices},
+        info={"scan_number": 1, "monitor_sync": "bec", **requested_info},
+    )
+    with (
+        mock.patch.object(manager.connector, "get", return_value=None),
+        mock.patch.object(manager, "write_file") as write,
+    ):
+        manager.update_scan_storage_with_status(closed_msg)
+        write.assert_called_once_with("scan_id")
+
+
+@pytest.mark.parametrize("first_update", ["status", "segment"])
+def test_pending_notifications_are_flushed_when_storage_is_created(
+    file_writer_manager_mock, first_update
+):
+    manager = file_writer_manager_mock
+    baseline = messages.ScanBaselineMessage(
+        scan_id="scan_id", data={"samz": {"samz": {"value": 42, "timestamp": 1}}}
+    )
+    files = {
+        name: messages.FileMessage(file_path=f"/{name}.h5", done=True, successful=True)
+        for name in ["eiger", "pilatus"]
+    }
+    manager._scan_baseline_callback(
+        MessageObject(topic=MessageEndpoints.scan_baseline().endpoint, value=baseline)
+    )
+    for name, file_msg in files.items():
+        manager._file_reference_callback(
+            MessageObject(
+                topic=MessageEndpoints.public_file("scan_id", name).endpoint, value=file_msg
+            )
+        )
+    manager._pending_baseline_writes["other_scan"] = {"samz": {}}
+    manager._pending_file_references["other_scan"] = {"eiger": files["eiger"]}
+    assert manager.scan_storage == {}
+
+    with mock.patch.object(manager, "write_file") as write:
+        if first_update == "status":
+            manager.update_scan_storage_with_status(
+                messages.ScanStatusMessage(
+                    scan_id="scan_id",
+                    status="paused",
+                    info={"scan_number": 1},
+                    readout_priority={"baseline": ["samz"]},
+                )
+            )
+        else:
+            manager.insert_to_scan_storage(
+                messages.ScanMessage(
+                    scan_id="scan_id", point_id=0, data={}, metadata={"scan_number": 1}
+                )
+            )
+        write.assert_not_called()
+
+    storage = manager.scan_storage["scan_id"]
+    assert storage.baseline == baseline.data
+    assert storage.file_references == files
+    assert "scan_id" not in manager._pending_baseline_writes
+    assert "scan_id" not in manager._pending_file_references
+    assert manager._pending_baseline_writes == {"other_scan": {"samz": {}}}
+    assert manager._pending_file_references == {"other_scan": {"eiger": files["eiger"]}}
+
+
+def test_baseline_before_first_status_allows_closed_scan_to_finalize(file_writer_manager_mock):
+    manager = file_writer_manager_mock
+    baseline = messages.ScanBaselineMessage(scan_id="scan_id", data={"samz": {}})
+    manager._scan_baseline_callback(
+        MessageObject(topic=MessageEndpoints.scan_baseline().endpoint, value=baseline)
+    )
+    with mock.patch.object(manager, "write_file") as write:
+        manager.update_scan_storage_with_status(
+            messages.ScanStatusMessage(
+                scan_id="scan_id",
+                status="closed",
+                scan_type="software_triggered",
+                num_monitored_readouts=0,
+                readout_priority={"monitored": [], "baseline": ["samz"]},
+                info={"scan_number": 1, "monitor_sync": "bec", "baseline_readout_requested": True},
+            )
+        )
+        write.assert_called_once_with("scan_id")
+    assert manager.scan_storage["scan_id"].baseline == baseline.data
+    assert manager._pending_baseline_writes == {}
+
+
+def test_file_reference_updates_active_storage(file_writer_manager_mock):
+    manager = file_writer_manager_mock
+    manager.scan_storage["scan_id"] = ScanStorage(1, "scan_id")
+    first = messages.FileMessage(file_path="/eiger.h5", done=False, successful=False)
+    finished = messages.FileMessage(file_path="/eiger.h5", done=True, successful=True)
+    for file_msg in [first, finished]:
+        manager._file_reference_callback(
+            MessageObject(
+                topic=MessageEndpoints.public_file("scan_id", "eiger").endpoint, value=file_msg
+            )
+        )
+    assert manager.scan_storage["scan_id"].file_references == {"eiger": finished}
+    assert manager._pending_file_references == {}
+
+
+@pytest.mark.parametrize(
+    "name,is_master_file", [("master", False), ("master", True), ("other", True)]
+)
+def test_master_file_notifications_are_not_buffered(file_writer_manager_mock, name, is_master_file):
+    manager = file_writer_manager_mock
+    file_msg = messages.FileMessage(
+        file_path="/master.h5", done=True, successful=True, is_master_file=is_master_file
+    )
+    manager._file_reference_callback(
+        MessageObject(topic=MessageEndpoints.public_file("scan_id", name).endpoint, value=file_msg)
+    )
+    assert manager.scan_storage == {}
+    assert manager._pending_file_references == {}
+
+
+def test_finished_scan_does_not_buffer_late_notifications(
+    file_writer_manager_mock, scan_storage_mock
+):
+    manager = file_writer_manager_mock
+    manager.scan_storage["scan_id"] = scan_storage_mock
+    with (
+        mock.patch(
+            "bec_server.file_writer.file_writer_manager.get_full_path", return_value="scan.h5"
+        ),
+        mock.patch.object(manager.file_writer, "write"),
+    ):
+        manager.write_file("scan_id")
+
+    baseline = messages.ScanBaselineMessage(scan_id="scan_id", data={"samz": {}})
+    file_msg = messages.FileMessage(file_path="/eiger.h5", done=True, successful=True)
+    manager._scan_baseline_callback(
+        MessageObject(topic=MessageEndpoints.scan_baseline().endpoint, value=baseline)
+    )
+    manager._file_reference_callback(
+        MessageObject(
+            topic=MessageEndpoints.public_file("scan_id", "eiger").endpoint, value=file_msg
+        )
+    )
+    assert manager.scan_storage == {}
+    assert manager._pending_baseline_writes == {}
+    assert manager._pending_file_references == {}
+
+
+def test_subscriptions_buffer_early_notifications(file_writer_manager_mock, connected_connector):
+    manager = file_writer_manager_mock
+    received = threading.Event()
+    baseline = messages.ScanBaselineMessage(scan_id="scan_id", data={"samz": {}})
+    file_msg = messages.FileMessage(file_path="/eiger.h5", done=True, successful=True)
+
+    def file_reference_callback(msg):
+        manager._file_reference_callback(msg)
+        received.set()
+
+    connected_connector.register(
+        MessageEndpoints.scan_baseline(), cb=manager._scan_baseline_callback
+    )
+    connected_connector.register(
+        patterns=MessageEndpoints.public_file("*", "*"), cb=file_reference_callback
+    )
+    pipe = connected_connector.pipeline()
+    connected_connector.set_and_publish(MessageEndpoints.scan_baseline(), baseline, pipe=pipe)
+    connected_connector.set_and_publish(
+        MessageEndpoints.public_file("scan_id", "eiger"), file_msg, pipe=pipe
+    )
+    pipe.execute()
+    assert received.wait(5), "File reference subscription did not deliver its notification"
+    assert manager.scan_storage == {}
+    assert manager._pending_baseline_writes == {"scan_id": baseline.data}
+    assert manager._pending_file_references == {"scan_id": {"eiger": file_msg}}
+
+    with mock.patch.object(manager, "write_file") as write:
+        manager.update_scan_storage_with_status(
+            messages.ScanStatusMessage(
+                scan_id="scan_id",
+                status="closed",
+                scan_type="software_triggered",
+                num_monitored_readouts=0,
+                readout_priority={"monitored": [], "baseline": ["samz"]},
+                info={"scan_number": 1, "monitor_sync": "bec", "baseline_readout_requested": True},
+            )
+        )
+        write.assert_called_once_with("scan_id")
+    assert manager.scan_storage["scan_id"].baseline == baseline.data
+    assert manager.scan_storage["scan_id"].file_references == {"eiger": file_msg}
+    assert manager._pending_baseline_writes == {}
+    assert manager._pending_file_references == {}

{bec_server-4.1.0 → bec_server-4.1.2}/tests/tests_scan_server/scans/test_scan_actions.py

@@ -13,6 +13,7 @@ from bec_lib.messaging_services import NotificationMessageObject
 from bec_lib.tests.fixtures import dm_with_devices  # noqa: F401
 from bec_lib.tests.utils import ConnectorMock
 from bec_lib.utils.scan_utils import compose_cli_input_from_scan_info
+from bec_server.file_writer.file_writer_manager import ScanStorage
 from bec_server.scan_server.instruction_handler import InstructionHandler
 from bec_server.scan_server.scans.scan_base import ScanBase, ScanInfo, ScanType
 from bec_server.scan_server.scans.scan_status import ScanStatus
@@ -388,6 +389,7 @@ def test_build_scan_status_message(action_context):
     assert msg.dataset_number == 2
     assert msg.num_points == 3
     assert msg.num_monitored_readouts == 0
+    assert msg.info["baseline_readout_requested"] is False
     assert msg.scan_type == "software_triggered"
     assert msg.scan_parameters == {
         "exp_time": 0.1,
@@ -576,6 +578,27 @@ def test_read_actions_emit_expected_messages_and_point_ids(action_context):
     assert monitored_msg_2.metadata["device_instr_id"] == monitored_status_2._device_instr_id
 
 
+@pytest.mark.parametrize("read_baseline", [False, True])
+def test_baseline_request_in_scan_status_controls_file_readiness(action_context, read_baseline):
+    ctx = action_context()
+    _set_readout_priority(ctx, baseline=["samz"], monitored=["samx"])
+    ctx.actions.read_monitored_devices(wait=False)
+    assert ctx.scan.scan_info.baseline_readout_requested is False
+    if read_baseline:
+        ctx.actions.read_baseline_devices(wait=False)
+
+    closed_msg = ctx.actions._build_scan_status_message("closed")
+    assert closed_msg.info["baseline_readout_requested"] is read_baseline
+    storage = ScanStorage(1, closed_msg.scan_id)
+    storage.status_msg = closed_msg
+    storage.scan_finished = True
+    storage.num_monitored_readouts = 1
+    storage.append(0, {"samx": {}})
+    assert storage.ready_to_write() is (not read_baseline)
+    storage.baseline = {"samz": {}}
+    assert storage.ready_to_write() is True
+
+
 def test_empty_read_and_trigger_actions_return_done_status(action_context):
     ctx = action_context()
     ctx.scan.scan_info.metadata["RID"] = "rid-123"
@@ -589,6 +612,7 @@ def test_empty_read_and_trigger_actions_return_done_status(action_context):
     trigger_status = ctx.actions.trigger_all_devices(wait=False)
 
     assert baseline_status.done
+    assert ctx.scan.scan_info.baseline_readout_requested is False
     assert monitored_status.done
     assert trigger_status.done
     assert not _sent_device_instructions(ctx, "read")