diff --git a/pyproject.toml b/pyproject.toml index a6a9659..4291de8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ "blueapi>=1.18.0", "fastapi>=0.136.0", "pydantic>=2.13.2", + "tiled>=0.2.9", ] dynamic = ["version"] license.file = "LICENSE" diff --git a/src/daq_queuing_service/api/errors.py b/src/daq_queuing_service/api/errors.py index c8ef493..6f9b2cd 100644 --- a/src/daq_queuing_service/api/errors.py +++ b/src/daq_queuing_service/api/errors.py @@ -1,6 +1,10 @@ +from collections.abc import Awaitable, Callable +from typing import TypeVar + from fastapi import FastAPI, Request from fastapi.responses import JSONResponse +from daq_queuing_service.log import LOGGER from daq_queuing_service.plugins.converter import ConverterError, ValidateError from daq_queuing_service.task_queue.queue_utils import ( NegativePositionError, @@ -12,59 +16,56 @@ # pyright: reportUnusedFunction=false +E = TypeVar("E", bound=Exception) -def register_exception_handlers(app: FastAPI): - @app.exception_handler(TaskInProgressError) - async def task_in_progress_handler( - request: Request, exception: TaskInProgressError - ): - return JSONResponse( - status_code=409, - content={"error": "task_in_progress", "message": str(exception)}, - ) +Handler = Callable[[Request, E], Awaitable[JSONResponse]] - @app.exception_handler(TaskNotFoundError) - async def task_not_found_handler(request: Request, exception: TaskNotFoundError): - return JSONResponse( - status_code=404, - content={"error": "task_not_found", "message": str(exception)}, - ) - @app.exception_handler(TaskNotInQueueError) - async def task_not_in_queue_handler( - request: Request, exception: TaskNotInQueueError - ): +def make_exception_handler( + status_code: int, error_code: str +) -> Callable[[Request, Exception], Awaitable[JSONResponse]]: + async def handler(request: Request, exception: Exception): + LOGGER.exception("Error while handling request: %s", request) return JSONResponse( - status_code=409, - content={"error": "task_not_in_queue", "message": str(exception)}, + status_code=status_code, + content={"error": error_code, "message": str(exception)}, ) - @app.exception_handler(NegativePositionError) - async def negative_position_handler( - request: Request, exception: NegativePositionError - ): - return JSONResponse( - status_code=400, - content={"error": "negative_position", "message": str(exception)}, - ) + return handler - @app.exception_handler(QueueError) - async def queue_error_handler(request: Request, exception: QueueError): - return JSONResponse( - status_code=409, - content={"error": "queue_error", "message": str(exception)}, - ) - @app.exception_handler(ValidateError) - async def validation_error_handler(request: Request, exception: ValidateError): - return JSONResponse( - status_code=422, - content={"error": "validation_error", "message": str(exception)}, - ) +def register_exception_handlers(app: FastAPI): + app.add_exception_handler( + TaskInProgressError, + make_exception_handler(409, "task_in_progress"), + ) - @app.exception_handler(ConverterError) - async def converter_error_handler(request: Request, exception: ConverterError): - return JSONResponse( - status_code=422, - content={"error": "converter_error", "message": str(exception)}, - ) + app.add_exception_handler( + TaskNotFoundError, + make_exception_handler(404, "task_not_found"), + ) + + app.add_exception_handler( + TaskNotInQueueError, + make_exception_handler(409, "task_not_in_queue"), + ) + + app.add_exception_handler( + NegativePositionError, + make_exception_handler(400, "negative_position"), + ) + + app.add_exception_handler( + QueueError, + make_exception_handler(409, "queue_error"), + ) + + app.add_exception_handler( + ValidateError, + make_exception_handler(422, "validation_error"), + ) + + app.add_exception_handler( + ConverterError, + make_exception_handler(422, "converter_error"), + ) diff --git a/src/daq_queuing_service/plugins/i15_1/backgrounds.py b/src/daq_queuing_service/plugins/i15_1/backgrounds.py new file mode 100644 index 0000000..179b272 --- /dev/null +++ b/src/daq_queuing_service/plugins/i15_1/backgrounds.py @@ -0,0 +1,24 @@ +from typing import Literal + +from pydantic import BaseModel, ConfigDict + +# This should be generated from the json schema +# https://github.com/DiamondLightSource/daq-queuing-service/issues/78 +BACKGROUND_TYPES = Literal["air", "bs", "fq", "pi"] + + +class BackgroundInfo(BaseModel): + # Currently only room temperatures scans are supported + # https://github.com/DiamondLightSource/daq-queuing-service/issues/84 + model_config = ConfigDict(frozen=True) + bg_type: BACKGROUND_TYPES + + def add_tiled_id(self, tiled_id: str) -> "TiledBackground": + return TiledBackground( + bg_type=self.bg_type, + tiled_id=tiled_id, + ) + + +class TiledBackground(BackgroundInfo): + tiled_id: str 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 new file mode 100644 index 0000000..bd82c4c --- /dev/null +++ b/src/daq_queuing_service/plugins/i15_1/i15_1_converter.py @@ -0,0 +1,190 @@ +from typing import Any + +from blueapi.service.model import TaskRequest +from tiled.client import from_uri # type: ignore +from tiled.client.container import Container + +from daq_queuing_service.blueapi_interaction.blueapi_call import BlueapiCall +from daq_queuing_service.log import LOGGER +from daq_queuing_service.plugins.converter import Converter +from daq_queuing_service.plugins.i15_1.backgrounds import BackgroundInfo +from daq_queuing_service.plugins.i15_1.tiled_interaction import get_background_tiled_id +from daq_queuing_service.task_queue.task import ( + Experiment, + ExperimentDefinition, + Sample, + Task, + TaskWithPosition, +) + +BACKGROUND_SCAN = "Background" + + +class I151Converter(Converter): + def __init__(self): + self.tiled_client: Container = from_uri("https://tiled.diamond.ac.uk/api/v1") + + def pre_process( + self, + queue: list[Task], + history: list[TaskWithPosition], + call_history: list[BlueapiCall], + ) -> list[Task]: + return self._add_required_background_scans(queue) + + def construct_blueapi_calls( + self, + queue: list[TaskWithPosition], + history: list[TaskWithPosition], + call_history: list[BlueapiCall], + ) -> list[BlueapiCall]: + + call_list: list[BlueapiCall] = [] + + for task in queue: + match task.experiment: + case TaskRequest(): + call_list.append( + BlueapiCall( + task_request=task.experiment, parent_task_id=task.id + ) + ) + case Experiment(): + call_list.extend( + [ + BlueapiCall(task_request=b_api_task, parent_task_id=task.id) + for b_api_task in ( + self._construct_blueapi_tasks_from_experiment( + task.experiment + ) + ) + ] + ) + return call_list + + def _construct_blueapi_tasks_from_experiment( + self, + experiment: Experiment, + ) -> list[TaskRequest]: + sample_name = experiment.sample.name + # Assume sample name is of form test_8_1 to load from position 8 on puck 1 + _, position, puck = sample_name.split("_") + + # For air calibration scans, we need to not to robot load/unload. + # https://github.com/DiamondLightSource/daq-queuing-service/issues/83 + return [ + TaskRequest( + name="robot_load", + params={"puck": puck, "position": position}, + instrument_session=experiment.instrument_session, + ), + TaskRequest( + name="centre_sample", + params={ + "start_z": -20, + "end_z": 0, + "steps": 20, + "exposure_time": 0.01, + "metadata": { + "sample": experiment.sample, + # This will include tiled background scan info + "experiment_definition": experiment.experiment_definition, + }, + }, + instrument_session=experiment.instrument_session, + ), + TaskRequest( + name="robot_unload", + params={}, + instrument_session=experiment.instrument_session, + ), + ] + + def _add_required_background_scans(self, tasks: list[Task]) -> list[Task]: + """Adds background scan tasks to the queue. Backgrounds will be added directly + in front of the first task in the queue that requires them. + + Args: + tasks (list[Task]): Current list of tasks + + Returns: + list[Task]: New list of tasks including backgrounds + """ + LOGGER.info("Adding required background scans") + + # This can be made more robust https://github.com/DiamondLightSource/daq-queuing-service/issues/80 + new_tasks: list[Task] = [] + + for task in tasks: + experiment = task.experiment + if ( + isinstance(experiment, Experiment) + and experiment.name != BACKGROUND_SCAN + ): + instrument_session = experiment.instrument_session + backgrounds = self._get_required_backgrounds(experiment) + + for background in backgrounds: + if tiled_id := get_background_tiled_id( + self.tiled_client, + background, + instrument_session, + ): + self._add_tiled_background_to_md( + experiment.experiment_definition.data, tiled_id, background + ) + + else: + bg_experiment = self._construct_background_experiment( + background, instrument_session + ) + new_tasks.append(Task(experiment=bg_experiment)) + + new_tasks.append(task) + return self._remove_repeated_backgrounds(new_tasks) + + def _remove_repeated_backgrounds(self, tasks: list[Task]) -> list[Task]: + LOGGER.info("Removing repeated background scans") + new_tasks: list[Task] = [] + queued_background_experiments: list[Experiment] = [] + + for task in tasks: + if task.experiment.name != BACKGROUND_SCAN: + new_tasks.append(task) + elif task.experiment not in queued_background_experiments: + assert isinstance(task.experiment, Experiment) + queued_background_experiments.append(task.experiment) + new_tasks.append(task) + else: + LOGGER.debug(f"Removing repeated background scan: {task.experiment}") + return new_tasks + + def _get_required_backgrounds(self, experiment: Experiment) -> list[BackgroundInfo]: + # This should be fleshed out https://github.com/DiamondLightSource/daq-queuing-service/issues/79 + return [BackgroundInfo(bg_type="fq")] + + def _add_tiled_background_to_md( + self, params: dict[str, Any], tiled_id: str, background: BackgroundInfo + ): + LOGGER.debug("Adding background scan tiled info to metadata") + if metadata := params.get("metadata"): + if tiled_backgrounds := metadata.get("tiled_backgrounds"): + tiled_backgrounds[tiled_id] = background + else: + metadata["tiled_backgrounds"] = {tiled_id: background} + else: + params["metadata"] = {"tiled_backgrounds": {tiled_id: background}} + + def _construct_background_experiment( + self, background: BackgroundInfo, instrument_session: str + ) -> Experiment: + LOGGER.debug(f"Constructing experiment for background: {background}") + return Experiment( + name=BACKGROUND_SCAN, + instrument_session=instrument_session, + # Need to get sample info for test samples (air, empty capillary etc) + sample=Sample(name="fq_1_1", id="", data={}), + experiment_definition=ExperimentDefinition( + name="background_scan", id="", data={"background": background} + ), + ) diff --git a/src/daq_queuing_service/plugins/i15_1/tiled_interaction.py b/src/daq_queuing_service/plugins/i15_1/tiled_interaction.py new file mode 100644 index 0000000..5216857 --- /dev/null +++ b/src/daq_queuing_service/plugins/i15_1/tiled_interaction.py @@ -0,0 +1,58 @@ +from cachetools import TTLCache, cached +from tiled.client.container import Container +from tiled.queries import Eq + +from daq_queuing_service.log import LOGGER +from daq_queuing_service.plugins.i15_1.backgrounds import BackgroundInfo + +# Ignoring the following rules as the tiled client is poorly typed as scares the linter +# pyright: reportUnknownMemberType=false +# pyright: reportUnknownVariableType=false +# pyright: reportUnknownArgumentType=false +# pyright: reportUnknownLambdaType=false + +cache: TTLCache[tuple[BackgroundInfo, str], str | None] = TTLCache(maxsize=100, ttl=1) + + +def get_background_tiled_id( + tiled_client: Container, + required_background: BackgroundInfo, + instrument_session: str, +) -> str | None: + + @cached(cache) + def _get_background_tiled_id( + required_background: BackgroundInfo, instrument_session: str + ) -> str | None: + + result: Container = ( + tiled_client.search(Eq("start.instrument_session", instrument_session)) + .search(Eq("start.instrument", "i15-1")) + .search( + Eq( + "start.experiment_definition.metadata.background", + required_background.model_dump_json(), + ) + ) + ) + + if not len(result): + LOGGER.debug( + f"Found no scans in tiled matching background: {required_background}" + ) + return + + items = sorted( + ((key, value) for key, value in result.items()), + key=lambda item: item[1].metadata["start"]["time"], + ) + + # return the tiled ID + tiled_id = items[-1][0] + LOGGER.debug( + f"Found {len(items)} scans in tiled matching background: " + + f"{required_background}. Returning the first: {tiled_id}" + ) + return tiled_id + + return _get_background_tiled_id(required_background, instrument_session) diff --git a/src/daq_queuing_service/plugins/i15_1_converter.py b/src/daq_queuing_service/plugins/i15_1_converter.py deleted file mode 100644 index 26948e9..0000000 --- a/src/daq_queuing_service/plugins/i15_1_converter.py +++ /dev/null @@ -1,73 +0,0 @@ -from blueapi.service.model import TaskRequest - -from daq_queuing_service.blueapi_interaction.blueapi_call import BlueapiCall -from daq_queuing_service.plugins.converter import Converter -from daq_queuing_service.task_queue.task import Experiment, TaskWithPosition - - -class I151Converter(Converter): - def _construct_blueapi_tasks_from_experiment( - self, - experiment: Experiment, - ) -> list[TaskRequest]: - sample_name = experiment.sample.name - # Assume sample name is of form test_8_1 to load from position 8 on puck 1 - _, position, puck = sample_name.split("_") - - return [ - TaskRequest( - name="robot_load", - params={"puck": puck, "position": position}, - instrument_session=experiment.instrument_session, - ), - TaskRequest( - name="centre_sample", - params={ - "start_z": -20, - "end_z": 0, - "steps": 20, - "exposure_time": 0.01, - "metadata": { - "sample": experiment.sample, - "experiment_definition": experiment.experiment_definition, - }, - }, - instrument_session=experiment.instrument_session, - ), - TaskRequest( - name="robot_unload", - params={}, - instrument_session=experiment.instrument_session, - ), - ] - - def construct_blueapi_calls( - self, - queue: list[TaskWithPosition], - history: list[TaskWithPosition], - call_history: list[BlueapiCall], - ) -> list[BlueapiCall]: - - call_list: list[BlueapiCall] = [] - - for task in queue: - match task.experiment: - case TaskRequest(): - call_list.append( - BlueapiCall( - task_request=task.experiment, parent_task_id=task.id - ) - ) - case Experiment(): - call_list.extend( - [ - BlueapiCall(task_request=b_api_task, parent_task_id=task.id) - for b_api_task in ( - self._construct_blueapi_tasks_from_experiment( - task.experiment - ) - ) - ] - ) - - return call_list diff --git a/src/daq_queuing_service/task_queue/task.py b/src/daq_queuing_service/task_queue/task.py index 6303fa0..016a187 100644 --- a/src/daq_queuing_service/task_queue/task.py +++ b/src/daq_queuing_service/task_queue/task.py @@ -1,4 +1,3 @@ -from collections.abc import Mapping from enum import StrEnum from typing import Any, Self from uuid import uuid4 @@ -16,13 +15,13 @@ class Sample(BaseModel): name: str id: str - data: Mapping[str, Any] + data: dict[str, Any] class ExperimentDefinition(BaseModel): name: str id: str - data: Mapping[str, Any] + data: dict[str, Any] class Experiment(BaseModel): diff --git a/tests/test_data/i15_1/test_daq_queue_config.yaml b/tests/test_data/i15_1/test_daq_queue_config.yaml index c022ab4..80937ec 100644 --- a/tests/test_data/i15_1/test_daq_queue_config.yaml +++ b/tests/test_data/i15_1/test_daq_queue_config.yaml @@ -1,6 +1,6 @@ converter: - path: "daq_queuing_service.plugins.i15_1_converter" - name: "construct_i15_1_blueapi_call_list" + path: "daq_queuing_service.plugins.i15_1.i15_1_converter" + name: "I151Converter" blueapi: api: url: "http://localhost:8000" diff --git a/tests/unit_tests/plugins/i15-1/conftest.py b/tests/unit_tests/plugins/i15-1/conftest.py new file mode 100644 index 0000000..4ab8eb5 --- /dev/null +++ b/tests/unit_tests/plugins/i15-1/conftest.py @@ -0,0 +1,19 @@ +from unittest.mock import patch + +import pytest + +from daq_queuing_service.plugins.i15_1.tiled_interaction import cache + + +@pytest.fixture(autouse=True) +def clear_cache(): + yield + cache.clear() + + +@pytest.fixture(autouse=True) +def tiled_client(): + with patch( + "daq_queuing_service.plugins.i15_1.i15_1_converter.from_uri" + ) as mock_from_uri: + yield mock_from_uri.return_value diff --git a/tests/unit_tests/plugins/i15-1/test_get_background_tiled_id.py b/tests/unit_tests/plugins/i15-1/test_get_background_tiled_id.py new file mode 100644 index 0000000..00220c2 --- /dev/null +++ b/tests/unit_tests/plugins/i15-1/test_get_background_tiled_id.py @@ -0,0 +1,85 @@ +from unittest.mock import MagicMock + +import pytest +from tiled.queries import Eq + +from daq_queuing_service.plugins.i15_1.backgrounds import BackgroundInfo +from daq_queuing_service.plugins.i15_1.tiled_interaction import get_background_tiled_id + + +@pytest.fixture() +def mock_tiled_searches( + tiled_client: MagicMock, +) -> tuple[MagicMock, MagicMock, MagicMock]: + result_1 = MagicMock() + result_1.metadata = {"start": {"time": 1}} + result_2 = MagicMock() + result_2.metadata = {"start": {"time": 10}} + result_3 = MagicMock() + result_3.metadata = {"start": {"time": 2}} + + search_result_3 = MagicMock() + search_result_3.search = MagicMock( + return_value={ + "tiled_id_1": result_1, + "tiled_id_2": result_2, + "tiled_id_3": result_3, + } + ) + + search_result_2 = MagicMock() + search_result_2.search = MagicMock(return_value=search_result_3) + + tiled_client.search = MagicMock(return_value=search_result_2) + + return tiled_client, search_result_2, search_result_3 + + +def test_get_background_tiled_id_makes_expected_searches( + mock_tiled_searches: tuple[MagicMock, MagicMock, MagicMock], +): + client, search_2, search_3 = mock_tiled_searches + get_background_tiled_id( + client, + BackgroundInfo(bg_type="air"), + instrument_session="cm12345-1", + ) + client.search.assert_called_once_with( + Eq(key="start.instrument_session", value="cm12345-1") + ) + search_2.search.assert_called_once_with(Eq(key="start.instrument", value="i15-1")) + search_3.search.assert_called_once_with( + Eq( + key="start.experiment_definition.metadata.background", + value='{"bg_type":"air"}', + ) + ) + + +def test_get_background_tiled_returns_most_recent_valid_background( + mock_tiled_searches: tuple[MagicMock, MagicMock, MagicMock], +): + client, _, _ = mock_tiled_searches + assert ( + get_background_tiled_id( + client, + BackgroundInfo(bg_type="air"), + instrument_session="cm12345-1", + ) + == "tiled_id_2" + ) + + +def test_get_background_tiled_id_returns_none_if_no_matching_backgrounds_found( + mock_tiled_searches: tuple[MagicMock, MagicMock, MagicMock], +): + client, _, final_search = mock_tiled_searches + final_search.search.return_value = {} + assert ( + get_background_tiled_id( + client, + BackgroundInfo(bg_type="air"), + instrument_session="cm12345-1", + ) + is None + ) 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 new file mode 100644 index 0000000..282c368 --- /dev/null +++ b/tests/unit_tests/plugins/i15-1/test_i15_1_converter.py @@ -0,0 +1,364 @@ +from copy import deepcopy +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from blueapi.service.model import TaskRequest + +from daq_queuing_service.blueapi_interaction.blueapi_call import BlueapiCall +from daq_queuing_service.plugins.i15_1.backgrounds import BackgroundInfo +from daq_queuing_service.plugins.i15_1.i15_1_converter import I151Converter +from daq_queuing_service.task_queue.task import ( + Experiment, + ExperimentDefinition, + Sample, + Status, + Task, + TaskKind, + TaskWithPosition, +) + + +def assert_tasks_equal(task1: Task | TaskWithPosition, task2: Task | TaskWithPosition): + # Check two tasks are equal other than the generated UUID + copy1 = type(task1).model_validate(task1) + copy2 = type(task2).model_validate(task2) + copy1.id = copy2.id = "" + assert task1 == task2 + + +@pytest.fixture(autouse=True) +def background_found_in_tiled(): + with patch( + "daq_queuing_service.plugins.i15_1.i15_1_converter.get_background_tiled_id", + MagicMock(return_value="fake_tiled_id"), + ) as mock_get_background_tiled_id: + yield mock_get_background_tiled_id + + +@pytest.fixture() +def background_not_found_in_tiled(): + with patch( + "daq_queuing_service.plugins.i15_1.i15_1_converter.get_background_tiled_id", + MagicMock(return_value=None), + ) as mock_get_background_tiled_id: + yield mock_get_background_tiled_id + + +@pytest.fixture() +def tasks_and_calls( + tasks: list[Task], +) -> tuple[list[TaskWithPosition], list[BlueapiCall]]: + tasks_with_positions = [TaskWithPosition.from_task(task) for task in tasks] + calls: list[BlueapiCall] = [] + for task in tasks_with_positions: + assert isinstance(task.experiment, Experiment) + calls.extend( + [ + BlueapiCall( + task_request=task_request, + parent_task_id=task.id, + ) + for task_request in ( + I151Converter()._construct_blueapi_tasks_from_experiment( + task.experiment + ) + ) + ] + ) + return tasks_with_positions, calls + + +def test_given_sample_name_in_correct_format_then_correct_sample_loaded(): + experiment = Experiment( + name="test_experiment", + experiment_definition=ExperimentDefinition(name=" ", id="", data={}), + sample=Sample(name="test_8_1", id="", data={}), + instrument_session="cm12345-1", + ) + tasks = I151Converter()._construct_blueapi_tasks_from_experiment(experiment) + assert tasks[0].name == "robot_load" + assert tasks[0].params["position"] == "8" + assert tasks[0].params["puck"] == "1" + + +def test_sample_centre_uses_expected_params(): + experiment = Experiment( + name="test_experiment", + experiment_definition=ExperimentDefinition(name=" ", id="", data={}), + sample=Sample(name="test_8_1", id="", data={}), + instrument_session="cm12345-1", + ) + tasks = I151Converter()._construct_blueapi_tasks_from_experiment(experiment) + assert tasks[1].name == "centre_sample" + assert tasks[1].params == { + "start_z": -20, + "end_z": 0, + "steps": 20, + "exposure_time": 0.01, + "metadata": { + "experiment_definition": ExperimentDefinition(name=" ", id="", data={}), + "sample": Sample(name="test_8_1", id="", data={}), + }, + } + + +def test_session_and_number_of_tasks_per_experiment_is_expected(): + experiment = Experiment( + name="test_experiment", + experiment_definition=ExperimentDefinition(name=" ", id="", data={}), + sample=Sample(name="test_8_1", id="", data={}), + instrument_session="cm12345-1", + ) + tasks = I151Converter()._construct_blueapi_tasks_from_experiment(experiment) + assert len(tasks) == 3 + for task in tasks: + assert task.instrument_session == "cm12345-1" + + +def test_experiment_with_correct_experiment_type_are_converted(): + experiment = Experiment( + name="test_experiment", + experiment_definition=ExperimentDefinition( + name="run_full_collection", id="", data={} + ), + sample=Sample(name="test_8_1", id="", data={}), + instrument_session="cm12345-1", + ) + task = TaskWithPosition( + experiment=experiment, + id="1", + status=Status.QUEUED, + blueapi_calls=[], + position=None, + kind=TaskKind.EXPERIMENT, + ) + call_list = I151Converter().construct_blueapi_calls([task], [], []) + assert len(call_list) == 3 + + +def test_mix_of_experiments_with_correct_experiment_type_are_converted(): + good_experiment = Experiment( + name="test_experiment", + experiment_definition=ExperimentDefinition( + name="run_full_collection", id="", data={} + ), + sample=Sample(name="test_8_1", id="", data={}), + instrument_session="cm12345-1", + ) + good_task = TaskWithPosition( + experiment=good_experiment, + id="1", + status=Status.QUEUED, + blueapi_calls=[], + position=None, + kind=TaskKind.EXPERIMENT, + ) + + class BadExperiment: + instrument_session = "cm12345-1" + + bad_task = deepcopy(good_task) + bad_task.experiment = BadExperiment() # type: ignore + plan_task = deepcopy(good_task) + plan_task.experiment = TaskRequest(name="", instrument_session="") + plan_task.kind = TaskKind.PLAN + tasks = [good_task, bad_task, plan_task, good_task] + call_list = I151Converter().construct_blueapi_calls(tasks, [], []) + assert len(call_list) == 7 + + +def test_if_no_background_found_in_tiled_then_background_scan_added_to_tasks( + background_not_found_in_tiled: None, +): + experiment = Experiment( + name="test_experiment", + experiment_definition=ExperimentDefinition( + name="run_full_collection", id="", data={} + ), + sample=Sample(name="test_8_1", id="", data={}), + instrument_session="cm12345-1", + ) + task = Task( + experiment=experiment, + id="1", + ) + tasks = I151Converter().pre_process([task], [], []) + assert len(tasks) == 2 + tasks[0].id = "" + assert tasks[0].model_dump() == { + "experiment": { + "name": "Background", + "instrument_session": "cm12345-1", + "sample": {"name": "fq_1_1", "id": "", "data": {}}, + "experiment_definition": { + "name": "background_scan", + "id": "", + "data": {"background": {"bg_type": "fq"}}, + }, + }, + "id": "", + "blueapi_calls": [], + "status": Status.QUEUED, + "kind": TaskKind.EXPERIMENT, + } + + +def test_add_required_background_scans_does_not_add_the_same_background_twice( + tasks: list[Task], background_not_found_in_tiled: None +): + bg_1 = BackgroundInfo(bg_type="air") + bg_2 = BackgroundInfo(bg_type="bs") + bg_3 = BackgroundInfo(bg_type="fq") + + def fake_get_required_background(self: I151Converter, experiment: Experiment): + # Get the same background scans every other experiment + # Only one of each background should be added + if int(experiment.sample.id) % 2 == 0: + return [bg_1, bg_2] + else: + return [bg_3] + + assert len(tasks) == 5 + with patch( + "daq_queuing_service.plugins.i15_1.i15_1_converter.I151Converter._get_required_backgrounds", + fake_get_required_background, + ): + new_tasks = I151Converter()._add_required_background_scans(tasks) + + assert len(new_tasks) == 8 + + assert_tasks_equal( + new_tasks[0], + Task( + experiment=I151Converter()._construct_background_experiment( + bg_1, instrument_session="" + ), + ), + ) + assert_tasks_equal( + new_tasks[1], + Task( + experiment=I151Converter()._construct_background_experiment( + bg_2, instrument_session="" + ), + ), + ) + # This one placed before the task that requires it + assert_tasks_equal( + new_tasks[3], + Task( + experiment=I151Converter()._construct_background_experiment( + bg_3, instrument_session="" + ), + ), + ) + + +def test_same_experiment_in_different_instrument_sessions_will_add_background_in_each( + tasks: list[Task], background_not_found_in_tiled: None +): + tasks[1].experiment.instrument_session = "different" + + assert len(tasks) == 5 + + new_tasks = I151Converter()._add_required_background_scans(tasks) + + assert len(new_tasks) == 7 + new_tasks[0].id = "" + assert new_tasks[0].model_dump() == { + "experiment": { + "name": "Background", + "instrument_session": "", + "sample": {"name": "fq_1_1", "id": "", "data": {}}, + "experiment_definition": { + "name": "background_scan", + "id": "", + "data": {"background": {"bg_type": "fq"}}, + }, + }, + "id": "", + "blueapi_calls": [], + "status": Status.QUEUED, + "kind": TaskKind.EXPERIMENT, + } + new_tasks[2].id = "" + assert new_tasks[2].model_dump() == { + "experiment": { + "name": "Background", + "instrument_session": "different", + "sample": {"name": "fq_1_1", "id": "", "data": {}}, + "experiment_definition": { + "name": "background_scan", + "id": "", + "data": {"background": {"bg_type": "fq"}}, + }, + }, + "id": "", + "blueapi_calls": [], + "status": Status.QUEUED, + "kind": TaskKind.EXPERIMENT, + } + + +def test_add_required_background_scans_if_found_in_tiled_then_no_background_added( + tasks: list[Task], + background_found_in_tiled: None, +): + tasks_after = I151Converter()._add_required_background_scans(tasks) + assert tasks_after == tasks + + +@pytest.mark.parametrize( + "params, tiled_ids, backgrounds, expected_params", + [ + ( + {"sample": "my_sample"}, + ["tiled_id"], + [BackgroundInfo(bg_type="bs")], + { + "metadata": { + "tiled_backgrounds": {"tiled_id": BackgroundInfo(bg_type="bs")} + }, + "sample": "my_sample", + }, + ), + ( + {}, + ["tiled_id"], + [BackgroundInfo(bg_type="bs")], + { + "metadata": { + "tiled_backgrounds": {"tiled_id": BackgroundInfo(bg_type="bs")} + }, + }, + ), + ( + {"sample": "my_sample"}, + ["tiled_id_1", "tiled_id_2"], + [ + BackgroundInfo(bg_type="bs"), + BackgroundInfo(bg_type="air"), + ], + { + "metadata": { + "tiled_backgrounds": { + "tiled_id_1": BackgroundInfo(bg_type="bs"), + "tiled_id_2": BackgroundInfo(bg_type="air"), + } + }, + "sample": "my_sample", + }, + ), + ], +) +def test_add_tiled_background_to_md_adds_expected_metadata( + params: dict[str, Any], + tiled_ids: list[str], + backgrounds: list[BackgroundInfo], + expected_params: dict[str, Any], +): + for tiled_id, background in zip(tiled_ids, backgrounds, strict=True): + I151Converter()._add_tiled_background_to_md(params, tiled_id, background) + + assert params == expected_params diff --git a/tests/unit_tests/plugins/test_i15_1_converter.py b/tests/unit_tests/plugins/test_i15_1_converter.py deleted file mode 100644 index 9a78589..0000000 --- a/tests/unit_tests/plugins/test_i15_1_converter.py +++ /dev/null @@ -1,111 +0,0 @@ -from copy import deepcopy - -from blueapi.service.model import TaskRequest - -from daq_queuing_service.plugins.i15_1_converter import I151Converter -from daq_queuing_service.task_queue.task import ( - Experiment, - ExperimentDefinition, - Sample, - Status, - TaskKind, - TaskWithPosition, -) - - -def test_given_sample_name_in_correct_format_then_correct_sample_loaded(): - experiment = Experiment( - name="test_experiment", - experiment_definition=ExperimentDefinition(name=" ", id="", data={}), - sample=Sample(name="test_8_1", id="", data={}), - instrument_session="cm12345-1", - ) - tasks = I151Converter()._construct_blueapi_tasks_from_experiment(experiment) - assert tasks[0].name == "robot_load" - assert tasks[0].params["position"] == "8" - assert tasks[0].params["puck"] == "1" - - -def test_sample_centre_uses_expected_params(): - experiment = Experiment( - name="test_experiment", - experiment_definition=ExperimentDefinition(name=" ", id="", data={}), - sample=Sample(name="test_8_1", id="", data={}), - instrument_session="cm12345-1", - ) - tasks = I151Converter()._construct_blueapi_tasks_from_experiment(experiment) - assert tasks[1].name == "centre_sample" - assert tasks[1].params == { - "start_z": -20, - "end_z": 0, - "steps": 20, - "exposure_time": 0.01, - "metadata": { - "experiment_definition": ExperimentDefinition(name=" ", id="", data={}), - "sample": Sample(name="test_8_1", id="", data={}), - }, - } - - -def test_session_and_number_of_tasks_per_experiment_is_expected(): - experiment = Experiment( - name="test_experiment", - experiment_definition=ExperimentDefinition(name=" ", id="", data={}), - sample=Sample(name="test_8_1", id="", data={}), - instrument_session="cm12345-1", - ) - tasks = I151Converter()._construct_blueapi_tasks_from_experiment(experiment) - assert len(tasks) == 3 - for task in tasks: - assert task.instrument_session == "cm12345-1" - - -def test_experiment_with_correct_experiment_type_are_converted(): - experiment = Experiment( - name="test_experiment", - experiment_definition=ExperimentDefinition( - name="run_full_collection", id="", data={} - ), - sample=Sample(name="test_8_1", id="", data={}), - instrument_session="cm12345-1", - ) - task = TaskWithPosition( - experiment=experiment, - id="1", - status=Status.QUEUED, - blueapi_calls=[], - position=None, - kind=TaskKind.EXPERIMENT, - ) - call_list = I151Converter().construct_blueapi_calls([task], [], []) - assert len(call_list) == 3 - - -def test_mix_of_experiments_with_correct_experiment_type_are_converted(): - good_experiment = Experiment( - name="test_experiment", - experiment_definition=ExperimentDefinition( - name="run_full_collection", id="", data={} - ), - sample=Sample(name="test_8_1", id="", data={}), - instrument_session="cm12345-1", - ) - good_task = TaskWithPosition( - experiment=good_experiment, - id="1", - status=Status.QUEUED, - blueapi_calls=[], - position=None, - kind=TaskKind.EXPERIMENT, - ) - - class BadExperiment: ... - - bad_task = deepcopy(good_task) - bad_task.experiment = BadExperiment() # type: ignore - plan_task = deepcopy(good_task) - plan_task.experiment = TaskRequest(name="", instrument_session="") - plan_task.kind = TaskKind.PLAN - tasks = [good_task, bad_task, plan_task, good_task] - call_list = I151Converter().construct_blueapi_calls(tasks, [], []) - assert len(call_list) == 7 diff --git a/uv.lock b/uv.lock index 04d7067..8ded321 100644 --- a/uv.lock +++ b/uv.lock @@ -955,6 +955,7 @@ dependencies = [ { name = "blueapi" }, { name = "fastapi" }, { name = "pydantic" }, + { name = "tiled" }, ] [package.dev-dependencies] @@ -983,6 +984,7 @@ requires-dist = [ { name = "blueapi", specifier = ">=1.18.0" }, { name = "fastapi", specifier = ">=0.136.0" }, { name = "pydantic", specifier = ">=2.13.2" }, + { name = "tiled", specifier = ">=0.2.9" }, ] [package.metadata.requires-dev]