bec-ipython-client 3.138.0__tar.gz → 3.139.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (43) hide show
  1. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/PKG-INFO +1 -1
  2. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/callbacks/ipython_live_updates.py +54 -10
  3. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/callbacks/utils.py +10 -4
  4. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/signals.py +27 -4
  5. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/pyproject.toml +3 -1
  6. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/client_tests/test_ipython_live_updates.py +285 -0
  7. bec_ipython_client-3.139.1/tests/client_tests/test_signals.py +145 -0
  8. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/.gitignore +0 -0
  9. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/__init__.py +0 -0
  10. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/beamline_mixin.py +0 -0
  11. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/bec_magics.py +0 -0
  12. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/bec_startup.py +0 -0
  13. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/callbacks/__init__.py +0 -0
  14. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/callbacks/device_progress.py +0 -0
  15. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/callbacks/live_table.py +0 -0
  16. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/callbacks/move_device.py +0 -0
  17. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/high_level_interfaces/__init__.py +0 -0
  18. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/high_level_interfaces/bec_hli.py +0 -0
  19. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/high_level_interfaces/spec_hli.py +0 -0
  20. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/main.py +0 -0
  21. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/plugins/SLS/__init__.py +0 -0
  22. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/plugins/SLS/sls_info.py +0 -0
  23. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/plugins/XTreme/__init__.py +0 -0
  24. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/plugins/XTreme/x-treme.py +0 -0
  25. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/plugins/__init__.py +0 -0
  26. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/plugins/flomni/flomni_config.yaml +0 -0
  27. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/prettytable.py +0 -0
  28. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/bec_ipython_client/progressbar.py +0 -0
  29. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/demo.py +0 -0
  30. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/client_tests/conftest.py +0 -0
  31. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/client_tests/test_beamline_mixins.py +0 -0
  32. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/client_tests/test_bec_client.py +0 -0
  33. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/client_tests/test_device_progress.py +0 -0
  34. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/client_tests/test_live_table.py +0 -0
  35. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/client_tests/test_move_callback.py +0 -0
  36. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/client_tests/test_pretty_table.py +0 -0
  37. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/conftest.py +0 -0
  38. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/end-2-end/_ensure_requirements_container.py +0 -0
  39. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/end-2-end/test_actors_e2e.py +0 -0
  40. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/end-2-end/test_procedures_e2e.py +0 -0
  41. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/end-2-end/test_scans_e2e.py +0 -0
  42. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/end-2-end/test_scans_lib_e2e.py +0 -0
  43. {bec_ipython_client-3.138.0 → bec_ipython_client-3.139.1}/tests/end-2-end/test_scans_v4_lib_e2e.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: bec_ipython_client
3
- Version: 3.138.0
3
+ Version: 3.139.1
4
4
  Summary: BEC IPython client
5
5
  Project-URL: Bug Tracker, https://github.com/bec-project/bec/issues
6
6
  Project-URL: Homepage, https://github.com/bec-project/bec
@@ -2,6 +2,7 @@ from __future__ import annotations
2
2
 
3
3
  import collections
4
4
  import time
5
+ from contextvars import Token
5
6
  from typing import TYPE_CHECKING, Any
6
7
 
7
8
  from rich.console import Console
@@ -11,6 +12,7 @@ from rich.panel import Panel
11
12
  from bec_ipython_client.callbacks.device_progress import LiveUpdatesDeviceProgress
12
13
  from bec_lib.bec_errors import ScanInterruption, ScanRestart
13
14
  from bec_lib.logger import bec_logger
15
+ from bec_lib.request_context import ActiveRequestContext, active_request_context
14
16
 
15
17
  from .live_table import LiveUpdatesTable
16
18
  from .move_device import LiveUpdatesReadbackProgressbar
@@ -164,11 +166,15 @@ class IPythonLiveUpdates:
164
166
  def process_request(self, request: messages.ScanQueueMessage, callbacks: Any) -> None:
165
167
  """Process the request and report instructions."""
166
168
  # pylint: disable=protected-access
169
+ context_token: Token | None = None
167
170
  try:
168
171
  with self.client._sighandler:
169
172
  # pylint: disable=protected-access
170
173
  self._active_request = request
171
174
  self._user_callback = callbacks
175
+ request_id = request.metadata["RID"]
176
+ request_context = ActiveRequestContext(request_id=request_id)
177
+ context_token = active_request_context.set(request_context)
172
178
  scan_request = ScanRequestMixin(self.client, request.metadata["RID"])
173
179
  scan_request.wait()
174
180
 
@@ -180,9 +186,11 @@ class IPythonLiveUpdates:
180
186
  time.sleep(0.01)
181
187
 
182
188
  self._current_queue = queue = scan_request.scan_queue_request.queue
189
+ request_context.queue_status = queue.status
183
190
  self._request_block_id = req_id = self._active_request.metadata.get("RID")
184
191
 
185
- while queue.status not in ["COMPLETED", "ABORTED", "HALTED"]:
192
+ while queue.status not in ["COMPLETED", "ABORTED", "HALTED", "CANCELLED"]:
193
+ request_context.queue_status = queue.status
186
194
  if self._process_queue(queue, request, req_id):
187
195
  break
188
196
 
@@ -209,7 +217,17 @@ class IPythonLiveUpdates:
209
217
  self._wait_for_cleanup()
210
218
  self._reset(forced=True)
211
219
  raise scan_interr
220
+ except KeyboardInterrupt as exc:
221
+ self._stop_status_live()
222
+ if self.client._service_config.abort_on_ctrl_c and self._abort_pending_request():
223
+ self._wait_for_cleanup()
224
+ self._reset(forced=True)
225
+ raise ScanInterruption("User abort.") from exc
226
+ self._reset(forced=True)
227
+ raise
212
228
  finally:
229
+ if context_token is not None:
230
+ active_request_context.reset(context_token)
213
231
  self._stop_status_live()
214
232
 
215
233
  def _wait_for_cleanup(self):
@@ -222,7 +240,32 @@ class IPythonLiveUpdates:
222
240
  while self._element_in_queue():
223
241
  time.sleep(0.1)
224
242
  except KeyboardInterrupt:
225
- self.client.queue.request_scan_halt()
243
+ request_id = self._get_tracked_request_id()
244
+ if request_id is None:
245
+ self.client.queue.request_scan_halt()
246
+ return
247
+ self.client.queue.request_scan_halt(request_id=request_id)
248
+
249
+ def _get_tracked_request_id(self) -> str | None:
250
+ """Return the request id associated with the active live-updates request."""
251
+ request_context = active_request_context.get()
252
+ if request_context is not None:
253
+ return request_context.request_id
254
+ if self._active_request is None:
255
+ return None
256
+ return self._active_request.metadata.get("RID")
257
+
258
+ def _abort_pending_request(self) -> bool:
259
+ """Abort the pending queue item currently being tracked by this live update."""
260
+ if self._current_queue is None or self.client.queue is None:
261
+ return False
262
+ if self._current_queue.status != "PENDING":
263
+ return False
264
+ request_id = self._get_tracked_request_id()
265
+ if request_id is None:
266
+ return False
267
+ self.client.queue.request_scan_abortion(request_id=request_id)
268
+ return True
226
269
 
227
270
  def _element_in_queue(self) -> bool:
228
271
  if self.client.queue is None:
@@ -237,7 +280,7 @@ class IPythonLiveUpdates:
237
280
  return False
238
281
  if self._current_queue is None:
239
282
  return False
240
- return self._current_queue.queue_id == queue_info[0].queue_id
283
+ return any(queue_item.queue_id == self._current_queue.queue_id for queue_item in queue_info)
241
284
 
242
285
  def _process_queue(
243
286
  self, queue: QueueItem, request: messages.ScanQueueMessage, req_id: str
@@ -259,7 +302,7 @@ class IPythonLiveUpdates:
259
302
  if queue.scans:
260
303
  if queue.scans[-1] is not None and queue.scans[-1].status == "user_completed":
261
304
  return True
262
- if queue.scans[-1] is None and queue.status == "STOPPED":
305
+ if queue.scans[-1] is None and queue.status in ["STOPPED", "CANCELLED"]:
263
306
  raise ScanInterruption("Scan was stopped by the user.")
264
307
 
265
308
  if queue.queue_position is None:
@@ -310,19 +353,20 @@ class IPythonLiveUpdates:
310
353
 
311
354
  if self.client.queue is None or self.client.queue.queue_storage.current_scan_queue is None:
312
355
  return
313
- target_queue = self.client.queue.queue_storage.current_scan_queue.get(
314
- self.client.queue.get_default_scan_queue()
315
- )
316
- if target_queue is None:
356
+ target_queue = self.client.queue.get_default_scan_queue()
357
+ target_queue_status = self.client.queue.queue_storage.current_scan_queue.get(target_queue)
358
+ if target_queue_status is None:
317
359
  return
318
360
 
319
361
  queue_position = queue.queue_position
320
362
  if queue_position is None:
321
363
  return
322
364
 
323
- status = target_queue.status
365
+ status = target_queue_status.status
324
366
  if status == "LOCKED":
325
- lock_info = [f"{lock.identifier}: {lock.reason}\n" for lock in target_queue.locks]
367
+ lock_info = [
368
+ f"{lock.identifier}: {lock.reason}\n" for lock in target_queue_status.locks
369
+ ]
326
370
  message = (
327
371
  f"Scan is waiting for the lock to be released. Active locks: \n{''.join(lock_info)}"
328
372
  )
@@ -74,20 +74,26 @@ class LiveUpdatesBase(abc.ABC):
74
74
  # pylint: disable=protected-access
75
75
  if self.scan_queue_request is None:
76
76
  return
77
- msgs = self.scan_queue_request.queue.get_client_messages(only_asap=True)
77
+ queue = self.scan_queue_request.queue
78
+ if queue is None:
79
+ return
80
+ msgs = queue.get_client_messages(only_asap=True)
78
81
  if not msgs:
79
82
  return
80
83
  if self.bec.live_updates_config.print_client_messages is False:
81
84
  return
82
85
  for msg in msgs:
83
- print(self.scan_queue_request.queue.format_client_msg(msg))
86
+ print(queue.format_client_msg(msg))
84
87
 
85
88
  def _print_client_msgs_all(self):
86
89
  """Print summary of client messages"""
87
90
  # pylint: disable=protected-access
88
91
  if self.scan_queue_request is None:
89
92
  return
90
- msgs = self.scan_queue_request.queue.get_client_messages()
93
+ queue = self.scan_queue_request.queue
94
+ if queue is None:
95
+ return
96
+ msgs = queue.get_client_messages()
91
97
  if self.bec.live_updates_config.print_client_messages is False:
92
98
  return
93
99
  if not msgs:
@@ -97,7 +103,7 @@ class LiveUpdatesBase(abc.ABC):
97
103
  print("------------------------")
98
104
  # pylint: disable=protected-access
99
105
  for msg in msgs:
100
- print(self.scan_queue_request.queue.format_client_msg(msg))
106
+ print(queue.format_client_msg(msg))
101
107
  print("------------------------")
102
108
 
103
109
 
@@ -6,6 +6,7 @@ import time
6
6
  from typing import TYPE_CHECKING
7
7
 
8
8
  from bec_lib.bec_errors import ScanInterruption
9
+ from bec_lib.request_context import active_request_context
9
10
 
10
11
  if TYPE_CHECKING: # pragma: no cover
11
12
  from bec_lib.client import BECClient
@@ -100,14 +101,28 @@ class SigintHandler(SignalHandler):
100
101
  else:
101
102
  raise ValueError(f"Mode {self._operation_mode} not handled by SigintHandler")
102
103
 
104
+ @staticmethod
105
+ def _get_active_request_id() -> str | None:
106
+ request_context = active_request_context.get()
107
+ if request_context is None:
108
+ return None
109
+ return request_context.request_id
110
+
103
111
  def _procedure_mode(self):
104
112
  # Catch it here to only kill scans which were started here
105
113
  print("SIGINT received in procedure mode. Sending scan abort request and exiting.")
106
- self.bec.queue.request_scan_abortion()
114
+ self.bec.queue.request_scan_abortion(request_id=self._get_active_request_id())
107
115
  # Let the procedure worker shut itself down
108
116
  raise KeyboardInterrupt
109
117
 
110
118
  def _normal_mode(self):
119
+ request_context = active_request_context.get()
120
+ request_id = request_context.request_id if request_context is not None else None
121
+ if request_context is not None and request_context.queue_status == "PENDING":
122
+ # A locally tracked request is still queued behind someone else.
123
+ # Let IPythonLiveUpdates handle Ctrl-C so we only target that request.
124
+ raise KeyboardInterrupt
125
+
111
126
  current_scan = self.bec.queue.scan_storage.current_scan_info
112
127
  if not current_scan:
113
128
  raise KeyboardInterrupt
@@ -126,7 +141,9 @@ class SigintHandler(SignalHandler):
126
141
  print("It has been 10 seconds since the last SIGINT. Resetting SIGINT handler.")
127
142
 
128
143
  threading.Thread(
129
- target=self.bec.queue.request_scan_interruption, args=(True,), daemon=True
144
+ target=self.bec.queue.request_scan_interruption,
145
+ kwargs={"deferred_pause": True, "request_id": request_id},
146
+ daemon=True,
130
147
  ).start()
131
148
  print(
132
149
  "A 'deferred pause' has been requested. The "
@@ -141,10 +158,16 @@ class SigintHandler(SignalHandler):
141
158
  # - Ctrl-C twice within 10 seconds or a direct command (e.g. mv) -> hard pause
142
159
  if self.bec._service_config.abort_on_ctrl_c:
143
160
  print("The scan will be aborted.")
144
- threading.Thread(target=self.bec.queue.request_scan_abortion, daemon=True).start()
161
+ threading.Thread(
162
+ target=self.bec.queue.request_scan_abortion,
163
+ kwargs={"request_id": request_id},
164
+ daemon=True,
165
+ ).start()
145
166
  raise ScanInterruption("User abort.")
146
167
  print("A hard pause will be requested.")
147
168
  threading.Thread(
148
- target=self.bec.queue.request_scan_interruption, args=(False,), daemon=True
169
+ target=self.bec.queue.request_scan_interruption,
170
+ kwargs={"deferred_pause": False, "request_id": request_id},
171
+ daemon=True,
149
172
  ).start()
150
173
  raise ScanInterruption(PAUSE_MSG)
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "bec_ipython_client"
7
- version = "3.138.0"
7
+ version = "3.139.1"
8
8
  description = "BEC IPython client"
9
9
  requires-python = ">=3.11"
10
10
  classifiers = [
@@ -94,6 +94,8 @@ Homepage = "https://github.com/bec-project/bec"
94
94
 
95
95
 
96
96
 
97
+
98
+
97
99
 
98
100
 
99
101
 
@@ -6,6 +6,7 @@ from bec_ipython_client.callbacks.ipython_live_updates import IPythonLiveUpdates
6
6
  from bec_lib import messages
7
7
  from bec_lib.bec_errors import ScanInterruption, ScanRestart
8
8
  from bec_lib.queue_items import QueueItem
9
+ from bec_lib.request_context import ActiveRequestContext, active_request_context
9
10
 
10
11
 
11
12
  @pytest.fixture
@@ -146,6 +147,41 @@ def test_live_updates_process_queue_running(ipython_live_updates_with_mocked_liv
146
147
  assert res is True
147
148
 
148
149
 
150
+ def test_live_updates_process_queue_cancelled_pending_request_raises_interruption(bec_client_mock):
151
+ client = bec_client_mock
152
+ live_updates = IPythonLiveUpdates(client)
153
+ request_msg = messages.ScanQueueMessage(
154
+ scan_type="grid_scan",
155
+ parameter={"args": {"samx": (-5, 5, 3)}, "kwargs": {}},
156
+ queue="primary",
157
+ metadata={"RID": "something"},
158
+ )
159
+ request_block = messages.RequestBlock(
160
+ msg=request_msg,
161
+ RID="something",
162
+ report_instructions=[],
163
+ readout_priority={"monitored": ["samx"]},
164
+ is_scan=True,
165
+ scan_number=1,
166
+ scan_id=None,
167
+ )
168
+ queue = QueueItem(
169
+ scan_manager=client.queue,
170
+ queue_id="queue_id",
171
+ request_blocks=[request_block],
172
+ status="CANCELLED",
173
+ active_request_block=None,
174
+ scan_id=[None],
175
+ )
176
+
177
+ with (
178
+ mock.patch.object(queue, "_update_with_buffer"),
179
+ mock.patch("bec_lib.queue_items.QueueItem.queue_position", new_callable=mock.PropertyMock),
180
+ pytest.raises(ScanInterruption, match="stopped by the user"),
181
+ ):
182
+ live_updates._process_queue(queue, request_msg, "something")
183
+
184
+
149
185
  def test_process_request_repeats_on_ScanRestart_error(
150
186
  ipython_live_updates_with_mocked_live, queue_elements
151
187
  ):
@@ -173,6 +209,222 @@ def test_process_request_repeats_on_ScanRestart_error(
173
209
  assert live_updates._stop_status_live.call_count == 5
174
210
 
175
211
 
212
+ def test_abort_pending_request_requests_abortion_by_request_id(bec_client_mock, sample_request_msg):
213
+ live_updates = IPythonLiveUpdates(bec_client_mock)
214
+ live_updates._current_queue = mock.MagicMock(status="PENDING", scan_ids=["scan_id", None])
215
+ live_updates._active_request = sample_request_msg
216
+
217
+ with mock.patch.object(live_updates.client.queue, "request_scan_abortion") as request_abort:
218
+ assert live_updates._abort_pending_request() is True
219
+ request_abort.assert_called_once_with(request_id="something")
220
+
221
+
222
+ def test_get_tracked_request_id_uses_contextvar(bec_client_mock):
223
+ live_updates = IPythonLiveUpdates(bec_client_mock)
224
+ live_updates._active_request = messages.ScanQueueMessage(
225
+ scan_type="grid_scan",
226
+ parameter={"args": {"samx": (-5, 5, 3)}, "kwargs": {}},
227
+ queue="primary",
228
+ metadata={"RID": "fallback-request"},
229
+ )
230
+ token = active_request_context.set(
231
+ ActiveRequestContext(request_id="context-request", queue_status="PENDING")
232
+ )
233
+
234
+ try:
235
+ assert live_updates._get_tracked_request_id() == "context-request"
236
+ finally:
237
+ active_request_context.reset(token)
238
+
239
+
240
+ def test_abort_pending_request_falls_back_to_request_id(bec_client_mock, sample_request_msg):
241
+ live_updates = IPythonLiveUpdates(bec_client_mock)
242
+ live_updates._current_queue = mock.MagicMock(status="PENDING", scan_ids=[None])
243
+ live_updates._active_request = sample_request_msg
244
+
245
+ with mock.patch.object(live_updates.client.queue, "request_scan_abortion") as request_abort:
246
+ assert live_updates._abort_pending_request() is True
247
+ request_abort.assert_called_once_with(request_id="something")
248
+
249
+
250
+ def test_abort_pending_request_returns_false_when_not_pending(bec_client_mock):
251
+ live_updates = IPythonLiveUpdates(bec_client_mock)
252
+ live_updates._current_queue = mock.MagicMock(status="RUNNING", scan_ids=["scan_id"])
253
+
254
+ with mock.patch.object(live_updates.client.queue, "request_scan_abortion") as request_abort:
255
+ assert live_updates._abort_pending_request() is False
256
+ request_abort.assert_not_called()
257
+
258
+
259
+ def test_wait_for_cleanup_returns_immediately_when_queue_item_gone(bec_client_mock):
260
+ live_updates = IPythonLiveUpdates(bec_client_mock)
261
+
262
+ with (
263
+ mock.patch.object(live_updates, "_stop_status_live") as stop_status_live,
264
+ mock.patch.object(live_updates, "_element_in_queue", return_value=False) as in_queue,
265
+ ):
266
+ live_updates._wait_for_cleanup()
267
+
268
+ stop_status_live.assert_called_once()
269
+ in_queue.assert_called_once()
270
+
271
+
272
+ def test_wait_for_cleanup_loops_until_queue_item_removed(bec_client_mock):
273
+ live_updates = IPythonLiveUpdates(bec_client_mock)
274
+
275
+ with (
276
+ mock.patch.object(
277
+ live_updates, "_element_in_queue", side_effect=[True, True, False]
278
+ ) as in_queue,
279
+ mock.patch("bec_ipython_client.callbacks.ipython_live_updates.time.sleep") as sleep,
280
+ ):
281
+ live_updates._wait_for_cleanup()
282
+
283
+ assert in_queue.call_count == 3
284
+ sleep.assert_called_once_with(0.1)
285
+
286
+
287
+ def test_wait_for_cleanup_keyboard_interrupt_requests_halt(bec_client_mock):
288
+ live_updates = IPythonLiveUpdates(bec_client_mock)
289
+ live_updates._active_request = messages.ScanQueueMessage(
290
+ scan_type="grid_scan",
291
+ parameter={"args": {"samx": (-5, 5, 3)}, "kwargs": {}},
292
+ queue="primary",
293
+ metadata={"RID": "pending-request"},
294
+ )
295
+
296
+ with (
297
+ mock.patch.object(live_updates, "_element_in_queue", side_effect=KeyboardInterrupt),
298
+ mock.patch.object(live_updates.client.queue, "request_scan_halt") as request_halt,
299
+ ):
300
+ live_updates._wait_for_cleanup()
301
+
302
+ request_halt.assert_called_once_with(request_id="pending-request")
303
+
304
+
305
+ def test_wait_for_cleanup_keyboard_interrupt_requests_halt_without_request_id(bec_client_mock):
306
+ live_updates = IPythonLiveUpdates(bec_client_mock)
307
+
308
+ with (
309
+ mock.patch.object(live_updates, "_element_in_queue", side_effect=KeyboardInterrupt),
310
+ mock.patch.object(live_updates.client.queue, "request_scan_halt") as request_halt,
311
+ ):
312
+ live_updates._wait_for_cleanup()
313
+
314
+ request_halt.assert_called_once_with()
315
+
316
+
317
+ @pytest.fixture
318
+ def process_request_keyboard_interrupt_setup(
319
+ ipython_live_updates_with_mocked_live, sample_request_msg
320
+ ):
321
+ live_updates, _ = ipython_live_updates_with_mocked_live
322
+ callbacks = mock.MagicMock()
323
+ live_updates.client._service_config = mock.MagicMock(abort_on_ctrl_c=True)
324
+
325
+ queue = mock.MagicMock()
326
+ queue.scan_ids = ["scan_id"]
327
+ queue.queue_id = "queue_id"
328
+
329
+ request_item = mock.MagicMock()
330
+ request_item.queue = queue
331
+ fake_scan_request = mock.MagicMock()
332
+ fake_scan_request.scan_queue_request = request_item
333
+ fake_scan_request.wait.return_value = None
334
+
335
+ live_updates.client._sighandler = mock.MagicMock()
336
+ live_updates.client._sighandler.__enter__.return_value = None
337
+ live_updates.client._sighandler.__exit__.return_value = None
338
+
339
+ return live_updates, callbacks, sample_request_msg, queue, fake_scan_request
340
+
341
+
342
+ def test_process_request_keyboard_interrupt_pending_request_raises_scan_interruption(
343
+ process_request_keyboard_interrupt_setup,
344
+ ):
345
+ live_updates, callbacks, sample_request_msg, queue, fake_scan_request = (
346
+ process_request_keyboard_interrupt_setup
347
+ )
348
+ queue.status = "PENDING"
349
+
350
+ with (
351
+ mock.patch(
352
+ "bec_ipython_client.callbacks.ipython_live_updates.ScanRequestMixin",
353
+ return_value=fake_scan_request,
354
+ ),
355
+ mock.patch.object(live_updates, "_process_queue", side_effect=KeyboardInterrupt()),
356
+ mock.patch.object(
357
+ live_updates, "_abort_pending_request", return_value=True
358
+ ) as abort_pending,
359
+ mock.patch.object(live_updates, "_wait_for_cleanup") as wait_for_cleanup,
360
+ mock.patch.object(live_updates, "_reset") as reset,
361
+ ):
362
+ with pytest.raises(ScanInterruption, match="User abort."):
363
+ live_updates.process_request(sample_request_msg, callbacks)
364
+
365
+ abort_pending.assert_called_once()
366
+ wait_for_cleanup.assert_called_once()
367
+ reset.assert_called_once_with(forced=True)
368
+
369
+
370
+ @pytest.mark.parametrize("queue_status", ["RUNNING", "LOCKED"])
371
+ def test_process_request_keyboard_interrupt_pending_request_aborts_local_request_by_id(
372
+ process_request_keyboard_interrupt_setup, queue_status
373
+ ):
374
+ live_updates, callbacks, sample_request_msg, queue, fake_scan_request = (
375
+ process_request_keyboard_interrupt_setup
376
+ )
377
+ queue.status = "PENDING"
378
+ live_updates.client.queue.queue_storage.current_scan_queue = {
379
+ "primary": messages.ScanQueueStatus(info=[], status=queue_status)
380
+ }
381
+
382
+ with (
383
+ mock.patch(
384
+ "bec_ipython_client.callbacks.ipython_live_updates.ScanRequestMixin",
385
+ return_value=fake_scan_request,
386
+ ),
387
+ mock.patch.object(live_updates, "_process_queue", side_effect=KeyboardInterrupt()),
388
+ mock.patch.object(live_updates.client.queue, "request_scan_abortion") as request_abort,
389
+ mock.patch.object(live_updates, "_wait_for_cleanup") as wait_for_cleanup,
390
+ mock.patch.object(live_updates, "_reset") as reset,
391
+ ):
392
+ with pytest.raises(ScanInterruption, match="User abort."):
393
+ live_updates.process_request(sample_request_msg, callbacks)
394
+
395
+ request_abort.assert_called_once_with(request_id="something")
396
+ wait_for_cleanup.assert_called_once()
397
+ reset.assert_called_once_with(forced=True)
398
+
399
+
400
+ def test_process_request_keyboard_interrupt_non_pending_re_raises(
401
+ process_request_keyboard_interrupt_setup,
402
+ ):
403
+ live_updates, callbacks, sample_request_msg, queue, fake_scan_request = (
404
+ process_request_keyboard_interrupt_setup
405
+ )
406
+ queue.status = "RUNNING"
407
+
408
+ with (
409
+ mock.patch(
410
+ "bec_ipython_client.callbacks.ipython_live_updates.ScanRequestMixin",
411
+ return_value=fake_scan_request,
412
+ ),
413
+ mock.patch.object(live_updates, "_process_queue", side_effect=KeyboardInterrupt()),
414
+ mock.patch.object(
415
+ live_updates, "_abort_pending_request", return_value=False
416
+ ) as abort_pending,
417
+ mock.patch.object(live_updates, "_wait_for_cleanup") as wait_for_cleanup,
418
+ mock.patch.object(live_updates, "_reset") as reset,
419
+ ):
420
+ with pytest.raises(KeyboardInterrupt):
421
+ live_updates.process_request(sample_request_msg, callbacks)
422
+
423
+ abort_pending.assert_called_once()
424
+ wait_for_cleanup.assert_not_called()
425
+ reset.assert_called_once_with(forced=True)
426
+
427
+
176
428
  @pytest.mark.timeout(20)
177
429
  def test_live_updates_process_queue_without_status(bec_client_mock, queue_elements):
178
430
  client = bec_client_mock
@@ -382,6 +634,39 @@ def test_element_in_queue_queue_id_in_info(bec_client_mock, sample_request_block
382
634
  assert live_updates._element_in_queue() is True
383
635
 
384
636
 
637
+ @pytest.mark.timeout(20)
638
+ def test_element_in_queue_queue_id_in_later_position(bec_client_mock, sample_request_block):
639
+ client = bec_client_mock
640
+ live_updates = IPythonLiveUpdates(client)
641
+
642
+ current_queue = mock.MagicMock()
643
+ current_queue.queue_id = "my_queue_id"
644
+ live_updates._current_queue = current_queue
645
+
646
+ first_entry = messages.QueueInfoEntry(
647
+ queue_id="different_queue_id",
648
+ scan_id=["scan_id_1"],
649
+ is_scan=[True],
650
+ request_blocks=[sample_request_block],
651
+ scan_number=[1],
652
+ status="RUNNING",
653
+ active_request_block=None,
654
+ )
655
+ second_entry = messages.QueueInfoEntry(
656
+ queue_id="my_queue_id",
657
+ scan_id=["scan_id_2"],
658
+ is_scan=[True],
659
+ request_blocks=[sample_request_block],
660
+ scan_number=[2],
661
+ status="PENDING",
662
+ active_request_block=None,
663
+ )
664
+ scan_queue_status = messages.ScanQueueStatus(info=[first_entry, second_entry], status="LOCKED")
665
+ client.queue.queue_storage.current_scan_queue = {"primary": scan_queue_status}
666
+
667
+ assert live_updates._element_in_queue() is True
668
+
669
+
385
670
  @pytest.mark.timeout(20)
386
671
  def test_process_pending_queue_element_locked_queue(
387
672
  ipython_live_updates_with_mocked_live, queue_elements
@@ -0,0 +1,145 @@
1
+ from __future__ import annotations
2
+
3
+ import threading
4
+ from unittest import mock
5
+
6
+ import pytest
7
+
8
+ from bec_ipython_client.signals import SigintHandler
9
+ from bec_lib.bec_errors import ScanInterruption
10
+ from bec_lib.request_context import ActiveRequestContext, active_request_context
11
+
12
+
13
+ @pytest.fixture
14
+ def bec_with_pending_live_request():
15
+ bec = mock.MagicMock()
16
+ bec._service_config = mock.MagicMock(abort_on_ctrl_c=True)
17
+ return bec
18
+
19
+
20
+ def test_sigint_handler_raises_keyboard_interrupt_for_pending_live_request(
21
+ bec_with_pending_live_request,
22
+ ):
23
+ handler = SigintHandler(bec_with_pending_live_request)
24
+ token = active_request_context.set(
25
+ ActiveRequestContext(request_id="pending-request", queue_status="PENDING")
26
+ )
27
+
28
+ try:
29
+ with pytest.raises(KeyboardInterrupt):
30
+ handler._normal_mode()
31
+ finally:
32
+ active_request_context.reset(token)
33
+
34
+ bec_with_pending_live_request.queue.request_scan_interruption.assert_not_called()
35
+ bec_with_pending_live_request.queue.request_scan_abortion.assert_not_called()
36
+
37
+
38
+ def test_sigint_handler_pending_live_request_does_not_abort_running_scan_from_other_client():
39
+ bec = mock.MagicMock()
40
+ bec._service_config = mock.MagicMock(abort_on_ctrl_c=True)
41
+ bec.queue.scan_storage.current_scan_info = mock.MagicMock(
42
+ status="RUNNING",
43
+ is_scan=[True],
44
+ request_blocks=[mock.MagicMock(RID="other-client-request")],
45
+ )
46
+ handler = SigintHandler(bec)
47
+ token = active_request_context.set(
48
+ ActiveRequestContext(request_id="local-pending-request", queue_status="PENDING")
49
+ )
50
+
51
+ try:
52
+ with pytest.raises(KeyboardInterrupt):
53
+ handler._normal_mode()
54
+ finally:
55
+ active_request_context.reset(token)
56
+
57
+ bec.queue.request_scan_interruption.assert_not_called()
58
+ bec.queue.request_scan_abortion.assert_not_called()
59
+
60
+
61
+ def test_sigint_handler_requests_deferred_pause_for_running_scan():
62
+ bec = mock.MagicMock()
63
+ bec._service_config = mock.MagicMock(abort_on_ctrl_c=True)
64
+ bec._live_updates = None
65
+ bec.queue.scan_storage.current_scan_info = mock.MagicMock(status="RUNNING", is_scan=[True])
66
+
67
+ handler = SigintHandler(bec)
68
+ token = active_request_context.set(
69
+ ActiveRequestContext(request_id="running-request", queue_status="RUNNING")
70
+ )
71
+
72
+ try:
73
+ with mock.patch.object(threading, "Thread") as thread_cls:
74
+ handler._normal_mode()
75
+ finally:
76
+ active_request_context.reset(token)
77
+
78
+ thread_cls.assert_called_once_with(
79
+ target=bec.queue.request_scan_interruption,
80
+ kwargs={"deferred_pause": True, "request_id": "running-request"},
81
+ daemon=True,
82
+ )
83
+ thread_cls.return_value.start.assert_called_once_with()
84
+
85
+
86
+ def test_sigint_handler_requests_abort_for_running_scan_after_second_sigint():
87
+ bec = mock.MagicMock()
88
+ bec._service_config = mock.MagicMock(abort_on_ctrl_c=True)
89
+ bec._live_updates = None
90
+ bec.queue.scan_storage.current_scan_info = mock.MagicMock(
91
+ status="DEFERRED_PAUSE", is_scan=[True]
92
+ )
93
+
94
+ handler = SigintHandler(bec)
95
+ handler.last_sigint_time = 0
96
+ token = active_request_context.set(
97
+ ActiveRequestContext(request_id="running-request", queue_status="DEFERRED_PAUSE")
98
+ )
99
+
100
+ try:
101
+ with (
102
+ mock.patch("bec_ipython_client.signals.time.time", return_value=5),
103
+ mock.patch.object(threading, "Thread") as thread_cls,
104
+ pytest.raises(ScanInterruption, match="User abort."),
105
+ ):
106
+ handler._normal_mode()
107
+ finally:
108
+ active_request_context.reset(token)
109
+
110
+ thread_cls.assert_called_once_with(
111
+ target=bec.queue.request_scan_abortion,
112
+ kwargs={"request_id": "running-request"},
113
+ daemon=True,
114
+ )
115
+ thread_cls.return_value.start.assert_called_once_with()
116
+
117
+
118
+ def test_sigint_handler_without_active_or_pending_scan_reraises_keyboard_interrupt():
119
+ bec = mock.MagicMock()
120
+ bec._service_config = mock.MagicMock(abort_on_ctrl_c=True)
121
+ bec.queue.scan_storage.current_scan_info = None
122
+ handler = SigintHandler(bec)
123
+
124
+ with pytest.raises(KeyboardInterrupt):
125
+ handler._normal_mode()
126
+
127
+ bec.queue.request_scan_interruption.assert_not_called()
128
+ bec.queue.request_scan_abortion.assert_not_called()
129
+
130
+
131
+ def test_sigint_handler_without_request_context_passes_none_to_thread_target():
132
+ bec = mock.MagicMock()
133
+ bec._service_config = mock.MagicMock(abort_on_ctrl_c=True)
134
+ bec._live_updates = None
135
+ bec.queue.scan_storage.current_scan_info = mock.MagicMock(status="RUNNING", is_scan=[True])
136
+ handler = SigintHandler(bec)
137
+
138
+ with mock.patch.object(threading, "Thread") as thread_cls:
139
+ handler._normal_mode()
140
+
141
+ thread_cls.assert_called_once_with(
142
+ target=bec.queue.request_scan_interruption,
143
+ kwargs={"deferred_pause": True, "request_id": None},
144
+ daemon=True,
145
+ )