From c4385eb580aaff7ddb99341de9a4be39a2899f51 Mon Sep 17 00:00:00 2001 From: Jacob Williamson Date: Fri, 6 Nov 2026 10:39:01 +0000 Subject: [PATCH 1/6] Add rules to tasks --- src/daq_queuing_service/task_queue/queue.py | 54 +++++++++++++++---- .../task_queue/queue_utils.py | 33 ++++++++++++ src/daq_queuing_service/task_queue/task.py | 14 ++++- .../plugins/i15_1/test_i15_1_converter.py | 5 ++ 4 files changed, 93 insertions(+), 13 deletions(-) diff --git a/src/daq_queuing_service/task_queue/queue.py b/src/daq_queuing_service/task_queue/queue.py index 598bc4fd..c0f7aa04 100644 --- a/src/daq_queuing_service/task_queue/queue.py +++ b/src/daq_queuing_service/task_queue/queue.py @@ -37,8 +37,10 @@ TaskNotClaimedError, TaskNotFoundError, TaskNotInQueueError, + TaskNotMovableError, + validate_task_position, ) -from daq_queuing_service.task_queue.task import Status, Task, TaskWithPosition +from daq_queuing_service.task_queue.task import Status, Task, TaskID, TaskWithPosition class TaskRegistry(dict[str, Task]): @@ -65,10 +67,10 @@ class QueueContents(TypedDict): tasks (dict[str, Task]): A dict containing all the tasks contained in the queue, including history, with the task ID as the key. - queue (list[str]): A list of the task IDs of the tasks in the queue. The order of + queue (list[TaskID]): A list of the task IDs of the tasks in the queue. The order of the list represents the order of the queue. - history (list[str]): A list of the task IDs of any completed tasks. + history (list[TaskID]): A list of the task IDs of any completed tasks. call_queue: (list[BlueapiCall]): A list of blueapi calls that will be executed in order to complete the queued tasks. @@ -77,8 +79,8 @@ class QueueContents(TypedDict): """ tasks: dict[str, Task] - queue: list[str] - history: list[str] + queue: list[TaskID] + history: list[TaskID] call_queue: list[BlueapiCall] call_history: list[BlueapiCall] @@ -135,8 +137,8 @@ def __init__( ): self.instrument = instrument self._tasks: TaskRegistry = TaskRegistry() - self._queue: list[str] = [] - self._history: list[str] = [] + self._queue: list[TaskID] = [] + self._history: list[TaskID] = [] self._call_queue: list[BlueapiCall] = [] self._call_history: list[BlueapiCall] = [] self._save_contents() @@ -204,6 +206,9 @@ def _sync(self): ] self._queue.extend(task.id for task in new_tasks) + # Ensure new task positions don't break any rules + self._validate_queue_positions() + # This is needed because pre_process could remove tasks self._remove_tasks_from_registry(list(self._tasks.keys())) @@ -503,8 +508,8 @@ async def move_task(self, task_id: str, position: int) -> int: int: The new position of the task (may be different to what was requested) """ async with self._modifying: - self._validate_tasks_for_move_or_deletion([task_id]) position = self._get_valid_position(position) + self._validate_task_for_move(task_id, position) self._remove_tasks_from_queue([task_id]) self._queue[position:position] = [task_id] new_position = self._queue.index(task_id) @@ -649,7 +654,7 @@ def _cancel_tasks(self, task_ids: Sequence[str]) -> list[TaskWithPosition]: task.cancel() return [TaskWithPosition.from_task(task) for task in tasks] - def _remove_tasks_from_queue(self, task_ids: list[str]) -> list[str]: + def _remove_tasks_from_queue(self, task_ids: list[TaskID]) -> list[TaskID]: # Only removes tasks in the queue (not history or registry) def should_be_removed(task_id: str): return ( @@ -662,7 +667,7 @@ def should_be_removed(task_id: str): return removed_ids - def _remove_tasks_from_registry(self, task_ids: list[str]) -> None: + def _remove_tasks_from_registry(self, task_ids: list[TaskID]) -> None: # Only removes tasks not present in the queue or history # So should remove tasks from queue/history before removing from registry def should_be_removed(task_id: str) -> bool: @@ -682,7 +687,7 @@ def should_be_removed(task_id: str) -> bool: } ) - def _validate_tasks_for_move_or_deletion(self, task_ids: list[str]): + def _validate_tasks_for_move_or_deletion(self, task_ids: list[TaskID]): for task_id in task_ids: task = self._tasks[task_id] if task_id not in self._queue: @@ -692,6 +697,31 @@ def _validate_tasks_for_move_or_deletion(self, task_ids: list[str]): f"Cannot move task '{task_id}', it is currently in progress!" ) + def _validate_task_for_move(self, task_id: TaskID, position: int): + task = self._tasks[task_id] + if not task.rules.movable: + raise TaskNotMovableError( + "Task isn't movable. See the following rules for task " + + f"{task_id}: {task.rules}" + ) + self._validate_task_in_position(task_id, position) + self._validate_tasks_for_move_or_deletion([task_id]) + + def _validate_task_in_position(self, task_id: TaskID, position: int): + task = self._tasks[task_id] + queue_without_task = [_id for _id in self._queue if _id != task_id] + queue_after_move = ( + queue_without_task[:position] + [task_id] + queue_without_task[position:] + ) + positions = {_id: i for i, _id in enumerate(queue_after_move)} + validate_task_position(task, positions) + + def _validate_queue_positions(self): + positions = {_id: i for i, _id in enumerate(self._queue)} + for task_id in self._queue: + task = self._tasks[task_id] + validate_task_position(task, positions) + def _get_queue(self) -> list[Task]: return [self._tasks[task_id] for task_id in self._queue] @@ -749,3 +779,5 @@ def _add_runtime_estimates_to_tasks(self): @cached_property def _tiled_client(self) -> Container: return get_tiled_client() + + def _check_rules(self): ... diff --git a/src/daq_queuing_service/task_queue/queue_utils.py b/src/daq_queuing_service/task_queue/queue_utils.py index c17491ee..cb7728aa 100644 --- a/src/daq_queuing_service/task_queue/queue_utils.py +++ b/src/daq_queuing_service/task_queue/queue_utils.py @@ -1,3 +1,6 @@ +from daq_queuing_service.task_queue.task import Task, TaskID + + class QueueError(Exception): pass @@ -24,3 +27,33 @@ class NegativePositionError(QueueError): class TaskNotClaimedError(QueueError): pass + + +class TaskNotMovableError(QueueError): + pass + + +class TaskPositionError(QueueError): + pass + + +def validate_task_position(task: Task, positions: dict[TaskID, int]): + task_position = positions[task.id] + + for must_be_before_id in task.rules.run_before: + must_be_before = positions.get(must_be_before_id) + if must_be_before is not None and task_position > must_be_before: + raise TaskPositionError( + f"Task cannot be in position {task_position} as it would be after" + + f" {must_be_before_id} in position {must_be_before}. See the " + + f"following rules for task {task.id}: {task.rules}" + ) + + for must_be_after_id in task.rules.run_after: + must_be_after = positions.get(must_be_after_id) + if must_be_after is not None and task_position < must_be_after: + raise TaskPositionError( + f"Task cannot be in position {task_position} as it would be before" + + f" {must_be_after_id} in position {must_be_after}. See the following " + + f"rules for task {task.id}: {task.rules}" + ) diff --git a/src/daq_queuing_service/task_queue/task.py b/src/daq_queuing_service/task_queue/task.py index 382cf43b..c6ed2606 100644 --- a/src/daq_queuing_service/task_queue/task.py +++ b/src/daq_queuing_service/task_queue/task.py @@ -1,5 +1,5 @@ from enum import StrEnum -from typing import Any, Self +from typing import Any, Self, TypeAlias from uuid import uuid4 from blueapi.service.model import TaskRequest @@ -47,6 +47,9 @@ def create_uuid_str() -> str: return str(uuid4()) +TaskID: TypeAlias = str + + class Status(StrEnum): QUEUED = "Queued" IN_PROGRESS = "In progress" @@ -60,12 +63,19 @@ class TaskKind(StrEnum): PLAN = "Plan" +class TaskRules(BaseModel): + movable: bool = True + run_before: set[TaskID] = Field(default_factory=set) + run_after: set[TaskID] = Field(default_factory=set) + + class Task(BaseModel): experiment: Experiment | TaskRequest - id: str = Field(default_factory=create_uuid_str) + id: TaskID = Field(default_factory=create_uuid_str) blueapi_calls: list[BlueapiCall] = Field(default_factory=lambda: []) user: User | None = None estimated_runtime: int | None = None + rules: TaskRules = Field(default_factory=TaskRules) _cancelled: bool = False def cancel(self): diff --git a/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py b/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py index 934eb708..ca9ac04b 100644 --- a/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py +++ b/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py @@ -211,6 +211,11 @@ def create_task_dict(task_id: str, experiment: dict[str, Any]) -> dict[str, Any] "user": None, "runtime": None, "estimated_runtime": None, + "rules": { + "movable": True, + "run_before": set(), + "run_after": set(), + }, } From 00359f748eea8fbbd78108154ec1687460dae5a3 Mon Sep 17 00:00:00 2001 From: Jacob Williamson Date: Fri, 6 Nov 2026 11:52:05 +0000 Subject: [PATCH 2/6] Add tests for position validation --- tests/unit_tests/test_queue_utils.py | 119 +++++++++++++++++++++++++++ 1 file changed, 119 insertions(+) create mode 100644 tests/unit_tests/test_queue_utils.py diff --git a/tests/unit_tests/test_queue_utils.py b/tests/unit_tests/test_queue_utils.py new file mode 100644 index 00000000..a7b90d4e --- /dev/null +++ b/tests/unit_tests/test_queue_utils.py @@ -0,0 +1,119 @@ +import pytest +from blueapi.service.model import TaskRequest + +from daq_queuing_service.task_queue.queue_utils import ( + TaskPositionError, + validate_task_position, +) +from daq_queuing_service.task_queue.task import Task, TaskID, TaskRules + + +@pytest.mark.parametrize( + "rules, queue", + [ + ( + TaskRules(run_after={"t0", "t1", "t2"}), + ["t0", "t1", "task_being_tested", "t2"], + ), + ( + TaskRules(run_after={"t0", "t5"}), + ["task_being_tested", "t0", "t1", "t2"], + ), + ( + TaskRules(run_before={"t2"}, run_after={"t3"}, movable=False), + ["t0", "t1", "task_being_tested", "t2", "t3"], + ), + ], +) +def test_validate_task_position_raises_error_if_task_too_early_in_queue( + rules: TaskRules, queue: list[TaskID] +): + positions = {task_id: i for i, task_id in enumerate(queue)} + task = Task( + id="task_being_tested", + rules=rules, + experiment=TaskRequest(name="", params={}, instrument_session=""), + ) + + with pytest.raises(TaskPositionError) as exc: + validate_task_position(task, positions) + + assert ( + f"Task cannot be in position {positions[task.id]} as it would be before" + in str(exc.value) + ) + + +@pytest.mark.parametrize( + "rules, queue", + [ + ( + TaskRules(run_before={"t0", "t1", "t2"}), + ["t0", "t1", "task_being_tested", "t2"], + ), + ( + TaskRules(run_before={"t0", "t5"}), + ["t0", "task_being_tested", "t1", "t2"], + ), + ( + TaskRules(run_before={"t2", "t1"}, run_after={"t0"}, movable=False), + ["t0", "t1", "task_being_tested", "t2", "t3"], + ), + ], +) +def test_validate_task_position_raises_error_if_task_too_late_in_queue( + rules: TaskRules, queue: list[TaskID] +): + positions = {task_id: i for i, task_id in enumerate(queue)} + task = Task( + id="task_being_tested", + rules=rules, + experiment=TaskRequest(name="", params={}, instrument_session=""), + ) + + with pytest.raises(TaskPositionError) as exc: + validate_task_position(task, positions) + + assert ( + f"Task cannot be in position {positions[task.id]} as it would be after" + in str(exc.value) + ) + + +@pytest.mark.parametrize( + "rules, queue", + [ + ( + TaskRules(run_before={"t2", "t3"}, run_after={"t0", "t1"}, movable=False), + ["t0", "t1", "task_being_tested", "t2", "t3"], + ), + ( + TaskRules( + run_before={"t2", "t3", "t4"}, run_after={"t0", "t1", "t5", "t6"} + ), + ["t0", "t1", "task_being_tested", "t2", "t3"], + ), + ( + TaskRules(run_before=set(), run_after=set()), + ["t0", "t1", "task_being_tested", "t2", "t3"], + ), + ( + TaskRules(run_before={"t1", "t2"}), + ["task_being_tested", "t0", "t1", "t2", "t3"], + ), + ( + TaskRules(run_after={"t1", "t2"}), + ["t0", "t1", "t2", "t3", "task_being_tested"], + ), + ], +) +def test_validate_task_position_does_not_raise_if_rules_met( + rules: TaskRules, queue: list[TaskID] +): + positions = {task_id: i for i, task_id in enumerate(queue)} + task = Task( + id="task_being_tested", + rules=rules, + experiment=TaskRequest(name="", params={}, instrument_session=""), + ) + validate_task_position(task, positions) From 020a26859569a1a2d310ea1b978f11f76993cede Mon Sep 17 00:00:00 2001 From: Jacob Williamson Date: Fri, 6 Nov 2026 12:57:00 +0000 Subject: [PATCH 3/6] Ensure you cannot have impossible rules --- src/daq_queuing_service/task_queue/task.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/daq_queuing_service/task_queue/task.py b/src/daq_queuing_service/task_queue/task.py index c6ed2606..c51dc1b7 100644 --- a/src/daq_queuing_service/task_queue/task.py +++ b/src/daq_queuing_service/task_queue/task.py @@ -3,7 +3,7 @@ from uuid import uuid4 from blueapi.service.model import TaskRequest -from pydantic import BaseModel, Field, computed_field +from pydantic import BaseModel, Field, computed_field, model_validator from daq_queuing_service.app.authentication import User from daq_queuing_service.external_interaction.blueapi.blueapi_call import ( @@ -68,6 +68,15 @@ class TaskRules(BaseModel): run_before: set[TaskID] = Field(default_factory=set) run_after: set[TaskID] = Field(default_factory=set) + @model_validator(mode="after") + def ensure_rules_do_not_contradict(self) -> Self: + if in_both := self.run_before.intersection(self.run_after): + raise ValueError( + f"{in_both} appear in both run_before and run_after. " + + "This is an impossible requirement." + ) + return self + class Task(BaseModel): experiment: Experiment | TaskRequest @@ -137,6 +146,7 @@ class TaskWithPosition(BaseModel): user: User | None estimated_runtime: int | None runtime: float | None + rules: TaskRules @classmethod def from_task(cls, task: Task, position: int | None = None) -> Self: From 00a03f48f442afa8104385c2a2ea1212a88f8a52 Mon Sep 17 00:00:00 2001 From: Jacob Williamson Date: Fri, 6 Nov 2026 12:57:19 +0000 Subject: [PATCH 4/6] Use rules in i15-1 plugin to ensure auxiliaries come first --- .../plugins/i15_1/i15_1_converter.py | 54 +++++++++++++------ 1 file changed, 38 insertions(+), 16 deletions(-) diff --git a/src/daq_queuing_service/plugins/i15_1/i15_1_converter.py b/src/daq_queuing_service/plugins/i15_1/i15_1_converter.py index 5819c899..3dc9e0f3 100644 --- a/src/daq_queuing_service/plugins/i15_1/i15_1_converter.py +++ b/src/daq_queuing_service/plugins/i15_1/i15_1_converter.py @@ -40,6 +40,8 @@ ExperimentDefinition, Sample, Task, + TaskID, + TaskRules, TaskWithPosition, ) @@ -52,9 +54,11 @@ ) -def _filter_auxiliary_scans(tasks: list[Task]) -> list[tuple[int, AuxiliaryScan]]: +def _filter_auxiliary_scans( + tasks: list[Task], +) -> list[tuple[int, TaskID, AuxiliaryScan]]: return [ - (i, AuxiliaryScan.from_experiment(task.experiment)) + (i, task.id, AuxiliaryScan.from_experiment(task.experiment)) for i, task in enumerate(tasks) if isinstance(task.experiment, Experiment) and is_auxiliary_experiment(task.experiment) @@ -63,9 +67,8 @@ def _filter_auxiliary_scans(tasks: list[Task]) -> list[tuple[int, AuxiliaryScan] class I151Converter(Converter): def __init__(self): - # First key is the ID of the task using the auxiliary scan self._tiled_auxiliary_scans: dict[ - str, dict[AuxiliaryScanType, TiledAuxiliary] + TaskID, dict[AuxiliaryScanType, TiledAuxiliary] ] = {} self._config_client = ConfigClient.from_url(CONFIG_SERVER_URL) @@ -144,7 +147,7 @@ def pause_on_fail_blueapi_call( return True def _construct_blueapi_tasks_from_experiment( - self, experiment: Experiment, task_id: str + self, experiment: Experiment, task_id: TaskID ) -> list[TaskRequest]: LOGGER.debug(f"Converting to blueapi calls, experiment = {experiment}") @@ -214,7 +217,7 @@ def _add_required_auxiliary_scans( required_auxiliary, current_task, new_tasks, - task.id, + task, instrument_session, ) @@ -226,7 +229,7 @@ def _ensure_auxiliary_in_queue_or_tiled( required_auxiliary: AuxiliaryScan, current_task: TaskWithPosition | None, new_tasks: list[Task], - task_id: str, + task: Task, instrument_session: str, ) -> list[Task]: if ( @@ -237,13 +240,19 @@ def _ensure_auxiliary_in_queue_or_tiled( required_auxiliary ) ): + task.rules.run_after.add(current_task.id) return new_tasks queued_auxiliary_scans = _filter_auxiliary_scans(new_tasks) - if any( - queued_auxiliary.is_suitable(required_auxiliary) - for _, queued_auxiliary in queued_auxiliary_scans + if suitable_task_id := next( + ( + task_id + for _, task_id, auxiliary_scan in queued_auxiliary_scans + if auxiliary_scan.is_suitable(required_auxiliary) + ), + None, ): + task.rules.run_after.add(suitable_task_id) return new_tasks if tiled_scan := get_suitable_tiled_scan( @@ -251,7 +260,7 @@ def _ensure_auxiliary_in_queue_or_tiled( required_auxiliary, self._queue_params.tiled_stale_time_s, ): - self._tiled_auxiliary_scans[task_id][tiled_scan.kind] = tiled_scan + self._tiled_auxiliary_scans[task.id][tiled_scan.kind] = tiled_scan return new_tasks LOGGER.info( @@ -261,6 +270,7 @@ def _ensure_auxiliary_in_queue_or_tiled( return self._add_or_replace_auxiliary_scan( required_auxiliary, new_tasks, + task, queued_auxiliary_scans, instrument_session, ) @@ -269,26 +279,38 @@ def _add_or_replace_auxiliary_scan( self, required_auxiliary: AuxiliaryScan, new_tasks: list[Task], - queued_auxiliary_scans: list[tuple[int, AuxiliaryScan]], + task: Task, + queued_auxiliary_scans: list[tuple[int, TaskID, AuxiliaryScan]], instrument_session: str, ) -> list[Task]: combined_auxiliary_scan = None index = None + aux_task_id = None - for i, auxiliary_scan in queued_auxiliary_scans: + for i, task_id, auxiliary_scan in queued_auxiliary_scans: if combined_auxiliary_scan := auxiliary_scan.attempt_to_combine_with( required_auxiliary ): + aux_task_id = task_id index = i break aux_experiment = self._construct_auxiliary_experiment( combined_auxiliary_scan or required_auxiliary, instrument_session ) - if index is None: - new_tasks.append(Task(experiment=aux_experiment)) + if index is None or aux_task_id is None: + new_tasks.append( + Task(experiment=aux_experiment, rules=TaskRules(movable=False)) + ) + aux_task_id = new_tasks[-1].id else: - new_tasks[index] = Task(experiment=aux_experiment) + new_tasks[index] = Task( + experiment=aux_experiment, + id=aux_task_id, + rules=TaskRules(movable=False), + ) + + task.rules.run_after.add(aux_task_id) return new_tasks def _get_required_auxiliary_scans( From 055e7fd5f9c5c580d108143fb63eb997729d635b Mon Sep 17 00:00:00 2001 From: Jacob Williamson Date: Fri, 6 Nov 2026 13:01:13 +0000 Subject: [PATCH 5/6] Fix tests --- .../plugins/i15_1/test_i15_1_converter.py | 37 ++++++++++++------- tests/unit_tests/test_api.py | 13 +++++++ tests/unit_tests/test_queue.py | 11 ++++++ 3 files changed, 47 insertions(+), 14 deletions(-) diff --git a/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py b/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py index ca9ac04b..2f5f4494 100644 --- a/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py +++ b/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py @@ -33,6 +33,7 @@ Status, Task, TaskKind, + TaskRules, TaskWithPosition, ) @@ -189,7 +190,8 @@ def make_background_task(capillary: STANDARD_CAPILLARY, time_per_pdf: int) -> Ta return Task( experiment=I151Converter()._construct_auxiliary_experiment( background, "cm12345-1" - ) + ), + rules=TaskRules(movable=False), ) @@ -201,7 +203,9 @@ def assert_tasks_equal(task1: Task | TaskWithPosition, task2: Task | TaskWithPos assert task1 == task2 -def create_task_dict(task_id: str, experiment: dict[str, Any]) -> dict[str, Any]: +def create_auxiliary_task_dict( + task_id: str, experiment: dict[str, Any] +) -> dict[str, Any]: return { "experiment": experiment, "id": task_id, @@ -212,7 +216,7 @@ def create_task_dict(task_id: str, experiment: dict[str, Any]) -> dict[str, Any] "runtime": None, "estimated_runtime": None, "rules": { - "movable": True, + "movable": False, "run_before": set(), "run_after": set(), }, @@ -222,7 +226,7 @@ def create_task_dict(task_id: str, experiment: dict[str, Any]) -> dict[str, Any] def air_task( task_id: str, instrument_session: str, time_per_pdf: float ) -> dict[str, Any]: - return create_task_dict( + return create_auxiliary_task_dict( task_id, { "name": "Air", @@ -254,7 +258,7 @@ def empty_capillary_task( "settle_time": 10, } ) - return create_task_dict( + return create_auxiliary_task_dict( task_id, { "name": f"Empty {capillary} - Empty Capillary", @@ -292,7 +296,7 @@ def standard_sample_task( contents: STANDARD_SAMPLE, standards_puck: StandardsPuck, ) -> dict[str, Any]: - return create_task_dict( + return create_auxiliary_task_dict( task_id, { "name": f"{contents} {capillary} - Standard Sample", @@ -486,6 +490,7 @@ def test_experiment_with_correct_experiment_type_are_converted(): user=None, runtime=None, estimated_runtime=100, + rules=TaskRules(), ) call_list = I151Converter().construct_blueapi_calls([task], [], []) assert len(call_list) == 6 @@ -660,6 +665,7 @@ def test_mix_of_experiments_with_correct_experiment_type_are_converted(): user=None, runtime=None, estimated_runtime=100, + rules=TaskRules(), ) class BadExperiment: @@ -744,6 +750,7 @@ def fake_get_required_auxiliary_scans(self: I151Converter, experiment: Experimen experiment=I151Converter()._construct_auxiliary_experiment( bg_1, instrument_session="cm12345-1" ), + rules=TaskRules(movable=False), ), ) assert_tasks_equal( @@ -752,6 +759,7 @@ def fake_get_required_auxiliary_scans(self: I151Converter, experiment: Experimen experiment=I151Converter()._construct_auxiliary_experiment( bg_2, instrument_session="cm12345-1" ), + rules=TaskRules(movable=False), ), ) # This one placed before the task that requires it @@ -761,6 +769,7 @@ def fake_get_required_auxiliary_scans(self: I151Converter, experiment: Experimen experiment=I151Converter()._construct_auxiliary_experiment( bg_3, instrument_session="cm12345-1" ), + rules=TaskRules(movable=False), ), ) @@ -882,9 +891,9 @@ def test_add_required_background_scans_if_found_in_tiled_then_no_background_adde async def test_queue_with_i15_1_converter_can_sync(queue_with_i15_1_plugin: TaskQueue): for i, task in enumerate(queue_with_i15_1_plugin._tasks.values()): task.experiment.instrument_session = str(i) - first_task = await queue_with_i15_1_plugin.get_task_by_position(0) - assert first_task - await queue_with_i15_1_plugin.move_task(first_task.id, 2) + a_task = await queue_with_i15_1_plugin.get_task_by_position(4) + assert a_task + await queue_with_i15_1_plugin.move_task(a_task.id, 3) call = await queue_with_i15_1_plugin.get_next_call_once_available() await queue_with_i15_1_plugin.return_call_to_queue(call) @@ -899,7 +908,7 @@ def test__ensure_background_in_queue_or_tiled_returns_if_suitable_already_queued ) new_tasks = [make_background_task("fq1.0", 25)] result = i15_1_converter._ensure_auxiliary_in_queue_or_tiled( - background, None, new_tasks, "task_id", "cm12345-1" + background, None, new_tasks, MagicMock(), "cm12345-1" ) assert result == new_tasks @@ -917,7 +926,7 @@ def test__ensure_background_in_queue_or_tiled_returns_if_current_task_is_suitabl background, TaskWithPosition.from_task(make_background_task("fq1.0", 25)), new_tasks, - "task_id", + MagicMock(), "cm12345-1", ) assert result == new_tasks @@ -937,7 +946,7 @@ def test__ensure_background_in_queue_or_tiled_modifies_queued_background_if_poss current_task.experiment.experiment_definition.name = "Not a background" new_tasks = [make_background_task("fq1.0", 10)] result = i15_1_converter._ensure_auxiliary_in_queue_or_tiled( - background, current_task, new_tasks, "task_id", "cm12345-1" + background, current_task, new_tasks, MagicMock(), "cm12345-1" ) assert len(result) == 1 assert_tasks_equal(result[0], make_background_task("fq1.0", 25)) @@ -953,7 +962,7 @@ def test__ensure_background_in_queue_or_tiled_adds_background_if_none_suitable_i time_per_pdf=25, ) result = i15_1_converter._ensure_auxiliary_in_queue_or_tiled( - background, None, [], "task_id", "cm12345-1" + background, None, [], MagicMock(), "cm12345-1" ) assert len(result) == 1 assert_tasks_equal(result[0], make_background_task("fq1.0", 25)) @@ -970,7 +979,7 @@ def test__ensure_background_in_queue_or_tiled_saves_tiled_info_if_exists( time_per_pdf=25, ) result = i15_1_converter._ensure_auxiliary_in_queue_or_tiled( - background, None, [], "task_id", "cm12345-1" + background, None, [], MagicMock(id="task_id"), "cm12345-1" ) assert len(result) == 0 assert i15_1_converter._tiled_auxiliary_scans == { diff --git a/tests/unit_tests/test_api.py b/tests/unit_tests/test_api.py index 1dd51775..e2bb5ce2 100644 --- a/tests/unit_tests/test_api.py +++ b/tests/unit_tests/test_api.py @@ -41,6 +41,7 @@ Status, Task, TaskKind, + TaskRules, TaskWithPosition, ) @@ -190,6 +191,7 @@ def test_get_queued_tasks_returns_queued_task(test_client: TestClient): "user": None, "runtime": None, "estimated_runtime": 100, + "rules": {"movable": True, "run_before": [], "run_after": []}, }, { "experiment": { @@ -230,6 +232,7 @@ def test_get_queued_tasks_returns_queued_task(test_client: TestClient): "user": None, "runtime": None, "estimated_runtime": 100, + "rules": {"movable": True, "run_before": [], "run_after": []}, }, { "experiment": { @@ -270,6 +273,7 @@ def test_get_queued_tasks_returns_queued_task(test_client: TestClient): "user": None, "runtime": None, "estimated_runtime": 100, + "rules": {"movable": True, "run_before": [], "run_after": []}, }, ] @@ -317,6 +321,7 @@ def test_get_queued_tasks_can_filter_by_task_status(test_client: TestClient): "user": None, "runtime": None, "estimated_runtime": 100, + "rules": {"movable": True, "run_before": [], "run_after": []}, } ] @@ -378,6 +383,7 @@ async def test_get_all_tasks_can_filter_by_task_status(test_client: TestClient): "user": None, "runtime": 59.0, "estimated_runtime": 100, + "rules": {"movable": True, "run_before": [], "run_after": []}, } ] @@ -440,6 +446,7 @@ async def test_add_tasks_to_queue_adds_to_queue_and_and_returns_task_ids( user=None, runtime=None, estimated_runtime=100, + rules=TaskRules(), ) @@ -725,6 +732,7 @@ async def test_cancel_tasks_removes_task_from_queue_and_returns_tasks( "user": None, "runtime": None, "estimated_runtime": 100, + "rules": {"movable": True, "run_before": [], "run_after": []}, }, { "experiment": { @@ -765,6 +773,7 @@ async def test_cancel_tasks_removes_task_from_queue_and_returns_tasks( "user": None, "runtime": None, "estimated_runtime": 100, + "rules": {"movable": True, "run_before": [], "run_after": []}, }, ] @@ -884,6 +893,7 @@ async def test_cancel_all_tasks_removes_all_queued_tasks_from_queue_and_returns_ "user": None, "runtime": None, "estimated_runtime": 100, + "rules": {"movable": True, "run_before": [], "run_after": []}, }, { "experiment": { @@ -931,6 +941,7 @@ async def test_cancel_all_tasks_removes_all_queued_tasks_from_queue_and_returns_ "user": None, "runtime": None, "estimated_runtime": 100, + "rules": {"movable": True, "run_before": [], "run_after": []}, }, ] @@ -984,6 +995,7 @@ def test_get_task_by_position_returns_expected_task(test_client: TestClient): "user": None, "runtime": None, "estimated_runtime": 100, + "rules": {"movable": True, "run_before": [], "run_after": []}, } @@ -1036,6 +1048,7 @@ def test_get_task_by_id_returns_expected_task(test_client: TestClient): "user": None, "runtime": None, "estimated_runtime": 100, + "rules": {"movable": True, "run_before": [], "run_after": []}, } diff --git a/tests/unit_tests/test_queue.py b/tests/unit_tests/test_queue.py index 366d4736..f7350433 100644 --- a/tests/unit_tests/test_queue.py +++ b/tests/unit_tests/test_queue.py @@ -38,6 +38,7 @@ Status, Task, TaskKind, + TaskRules, ) from .conftest import make_sample @@ -315,6 +316,7 @@ async def test_get_queue_only_returns_tasks_in_queue( user=None, runtime=None, estimated_runtime=100, + rules=TaskRules(), ), TaskWithPosition( experiment=Experiment( @@ -347,6 +349,7 @@ async def test_get_queue_only_returns_tasks_in_queue( user=None, runtime=None, estimated_runtime=100, + rules=TaskRules(), ), TaskWithPosition( experiment=Experiment( @@ -379,6 +382,7 @@ async def test_get_queue_only_returns_tasks_in_queue( user=None, runtime=None, estimated_runtime=100, + rules=TaskRules(), ), ] @@ -426,6 +430,7 @@ async def test_get_history_only_returns_tasks_in_history( user=None, runtime=59, estimated_runtime=100, + rules=TaskRules(), ), TaskWithPosition( experiment=Experiment( @@ -458,6 +463,7 @@ async def test_get_history_only_returns_tasks_in_history( user=None, runtime=59, estimated_runtime=100, + rules=TaskRules(), ), ] @@ -506,6 +512,7 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( user=None, runtime=59, estimated_runtime=100, + rules=TaskRules(), ), TaskWithPosition( experiment=Experiment( @@ -538,6 +545,7 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( user=None, runtime=59, estimated_runtime=100, + rules=TaskRules(), ), TaskWithPosition( experiment=Experiment( @@ -570,6 +578,7 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( user=None, runtime=None, estimated_runtime=100, + rules=TaskRules(), ), TaskWithPosition( experiment=Experiment( @@ -602,6 +611,7 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( user=None, runtime=None, estimated_runtime=100, + rules=TaskRules(), ), TaskWithPosition( experiment=Experiment( @@ -634,6 +644,7 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( user=None, runtime=None, estimated_runtime=100, + rules=TaskRules(), ), ] From a725a4a838e3c48457ad157403f2f972b90b2d01 Mon Sep 17 00:00:00 2001 From: Jacob Williamson Date: Fri, 6 Nov 2026 13:48:50 +0000 Subject: [PATCH 6/6] Add tests --- .../plugins/i15_1/test_i15_1_converter.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py b/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py index 2f5f4494..6d3a56f8 100644 --- a/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py +++ b/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py @@ -20,10 +20,15 @@ AuxiliaryScan, AuxiliaryScanType, TiledAuxiliary, + is_auxiliary_experiment, ) from daq_queuing_service.plugins.i15_1.i15_1_converter import I151Converter from daq_queuing_service.plugins.i15_1.plans import ScanType from daq_queuing_service.task_queue.queue import TaskQueue +from daq_queuing_service.task_queue.queue_utils import ( + TaskNotMovableError, + TaskPositionError, +) from daq_queuing_service.task_queue.task import ( Container, ContainerPosition, @@ -1117,3 +1122,24 @@ def test_no_wait_for_beam_queued_if_turned_off_in_queue_params( [TaskWithPosition.from_task(i15_1_tasks[0])], [], [] ) assert not any(call.task_request.name == "wait_for_beam" for call in blueapi_calls) + + +async def test_auxiliary_scans_cannot_be_moved(queue_with_i15_1_plugin: TaskQueue): + for i in range(len(queue_with_i15_1_plugin._queue)): + task = await queue_with_i15_1_plugin.get_task_by_position(i) + assert task and isinstance(task.experiment, Experiment) + if is_auxiliary_experiment(task.experiment): + with pytest.raises(TaskNotMovableError): + await queue_with_i15_1_plugin.move_task(task.id, 4) + + +async def test_scans_cannot_be_moved_before_auxiliaries( + queue_with_i15_1_plugin: TaskQueue, +): + for i in range(len(queue_with_i15_1_plugin._queue)): + task = await queue_with_i15_1_plugin.get_task_by_position(i) + assert task and isinstance(task.experiment, Experiment) + if not is_auxiliary_experiment(task.experiment): + for j in range(3): + with pytest.raises(TaskPositionError): + await queue_with_i15_1_plugin.move_task(task.id, j)