diff --git a/docs/reference/openapi.yaml b/docs/reference/openapi.yaml index e0cd8282d..d4ccda4bd 100644 --- a/docs/reference/openapi.yaml +++ b/docs/reference/openapi.yaml @@ -246,10 +246,8 @@ components: title: Name type: string params: - additionalProperties: true + $ref: '#/components/schemas/TaskParametersV2' description: Values for parameters to plan, if any - title: Params - type: object required: - name title: Task @@ -274,9 +272,22 @@ components: - message title: TaskError type: object - TaskRequest: + TaskParametersV2: additionalProperties: false - description: Request to run a task with related info + properties: + args: + items: {} + title: Args + type: array + kwargs: + additionalProperties: true + title: Kwargs + type: object + title: TaskParametersV2 + type: object + TaskRequestV1: + additionalProperties: false + description: Request to run a task with related info. Supports **kwargs only. properties: instrument_session: description: Instrument session associated with this task @@ -294,7 +305,27 @@ components: required: - name - instrument_session - title: TaskRequest + title: TaskRequestV1 + type: object + TaskRequestV2: + additionalProperties: false + description: Request to run a task with related info. Supports *args and **kwargs. + properties: + instrument_session: + description: Instrument session associated with this task + title: Instrument Session + type: string + name: + description: Name of plan to run + title: Name + type: string + params: + $ref: '#/components/schemas/TaskParametersV2' + description: Values for parameters to plan, if any + required: + - name + - instrument_session + title: TaskRequestV2 type: object TaskResponse: additionalProperties: false @@ -449,7 +480,7 @@ info: name: Apache 2.0 url: https://www.apache.org/licenses/LICENSE-2.0.html title: BlueAPI Control - version: 1.5.0 + version: 1.6.0 openapi: 3.1.0 paths: /api/v1/devices: @@ -641,13 +672,18 @@ paths: content: application/json: schema: - $ref: '#/components/schemas/TaskRequest' + anyOf: + - $ref: '#/components/schemas/TaskRequestV2' + - $ref: '#/components/schemas/TaskRequestV1' examples: - instrument_session: cm12345-1 name: count params: - detectors: - - x + args: [] + kwargs: + detectors: + - x + title: Task Request required: true responses: '201': @@ -1042,13 +1078,18 @@ paths: content: application/json: schema: - $ref: '#/components/schemas/TaskRequest' + anyOf: + - $ref: '#/components/schemas/TaskRequestV2' + - $ref: '#/components/schemas/TaskRequestV1' examples: - instrument_session: cm12345-1 name: count params: - detectors: - - x + args: [] + kwargs: + detectors: + - x + title: Task Request required: true responses: '201': diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py index c3f9bde13..070aeed73 100644 --- a/src/blueapi/cli/cli.py +++ b/src/blueapi/cli/cli.py @@ -4,7 +4,7 @@ import stat import sys import textwrap -from collections.abc import Callable +from collections.abc import Callable, Mapping from functools import wraps from pathlib import Path from pprint import pprint @@ -38,8 +38,13 @@ from blueapi.core import OTLP_EXPORT_ENABLED, DataEvent from blueapi.log import set_up_logging from blueapi.service.authentication import SessionCacheManager, SessionManager -from blueapi.service.model import DeviceResponse, PlanResponse, SourceInfo, TaskRequest -from blueapi.worker import ProgressEvent, WorkerEvent +from blueapi.service.model import ( + DeviceResponse, + PlanResponse, + SourceInfo, + TaskRequestV2, +) +from blueapi.worker import ProgressEvent, TaskParametersV2, WorkerEvent from blueapi.worker.event import TaskError, TaskResult from .scratch import setup_scratch @@ -47,8 +52,6 @@ LOGGER = logging.getLogger(__name__) -TaskParameters = dict[str, Any] - class ParametersType(ParamType): """CLI input parameter to accept a JSON object as an argument""" @@ -57,23 +60,36 @@ class ParametersType(ParamType): def convert( self, - value: str | dict[str, Any] | None, + value: str | Mapping[str, Any] | None, param: Parameter | None, ctx: Context | None, - ) -> TaskParameters: + ) -> TaskParametersV2: + if value is None: + return TaskParametersV2() + if isinstance(value, str): try: - params = json.loads(value) - if is_str_dict(params): - return params - self.fail("Parameters must be a JSON object with string keys") + value = json.loads(value) except json.JSONDecodeError as jde: - self.fail(f"Parameters are not valid JSON: {jde}") - else: - return super().convert(value, param, ctx) + self.fail(f"Parameters are not valid JSON: {jde}", param, ctx) + + if not isinstance(value, dict) or not all( + isinstance(key, str) for key in value + ): + self.fail( + "Parameters must be a JSON object with string keys", + param, + ctx, + ) + # V2 parameters have the shape {"args": [...], "kwargs": {...}}. + if value and set(value).issubset({"args", "kwargs"}): + return TaskParametersV2.model_validate(value) + + # Otherwise, interpret the object as V1 keyword arguments. + return TaskParametersV2.from_v1(value) -def is_str_dict(val: Any) -> TypeGuard[TaskParameters]: +def is_str_dict(val: Any) -> TypeGuard[TaskParametersV2]: return isinstance(val, dict) and all(isinstance(k, str) for k in val) @@ -363,7 +379,7 @@ def run_plan( foreground: bool, ws: bool, instrument_session: str, - parameters: TaskParameters, + parameters: TaskParametersV2, ) -> None: """Run a plan with parameters @@ -371,10 +387,11 @@ def run_plan( configuration or the experimental --ws flag is required. Without stomp configuration, '--bg' can be used to start a plan in the background. """ - client = cast(BlueapiClient, obj["client"]) - task = TaskRequest( - name=name, params=parameters, instrument_session=instrument_session + task = TaskRequestV2( + name=name, + params=parameters, + instrument_session=instrument_session, ) try: diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index e393da659..793611afb 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -28,7 +28,7 @@ PlanResponse, PythonEnvironmentResponse, SourceInfo, - TaskRequest, + TaskRequestV2, TaskResponse, TasksListResponse, WorkerTask, @@ -313,7 +313,7 @@ def get_active_task(self) -> WorkerTask: @start_as_current_span(TRACER, "request") def run_blocking( - self, request: TaskRequest, on_event: OnAnyEvent | None = None + self, request: TaskRequestV2, on_event: OnAnyEvent | None = None ) -> TaskStatus: for event in self._rest.run_blocking(request): if on_event is not None: @@ -335,7 +335,7 @@ def run_blocking( @start_as_current_span(TRACER, "task", "timeout") def run_task( self, - task: TaskRequest, + task: TaskRequestV2, on_event: OnAnyEvent | None = None, timeout: float | None = None, ) -> TaskStatus: @@ -403,7 +403,7 @@ def inner_on_event(event: AnyEvent, ctx: MessageContext) -> None: return complete.result(timeout=timeout) @start_as_current_span(TRACER, "task") - def create_and_start_task(self, task: TaskRequest) -> TaskResponse: + def create_and_start_task(self, task: TaskRequestV2) -> TaskResponse: """ Create a new task and instruct the worker to start it immediately. @@ -429,7 +429,7 @@ def create_and_start_task(self, task: TaskRequest) -> TaskResponse: @start_as_current_span(TRACER, "task") @deprecated("rest client") - def create_task(self, task: TaskRequest) -> TaskResponse: + def create_task(self, task: TaskRequestV2) -> TaskResponse: """ Create a new task, does not start execution diff --git a/src/blueapi/client/plan_cache.py b/src/blueapi/client/plan_cache.py index 69e64a1ab..210cf873b 100644 --- a/src/blueapi/client/plan_cache.py +++ b/src/blueapi/client/plan_cache.py @@ -1,7 +1,7 @@ import logging from typing import Any -from blueapi.service.model import PlanModel, TaskRequest +from blueapi.service.model import PlanModel, TaskParametersV2, TaskRequestV2 from blueapi.worker.event import TaskError, TaskResult, TaskStatus from .protocols import ClientProtocol @@ -26,8 +26,6 @@ def __init__(self, typ: str, message: str): class PlanCache: - """Collection of plans that can be accessed by name or attribute.""" - def __init__(self, client: ClientProtocol, plans: list[PlanModel]): self._cache = { model.name: Plan(name=model.name, model=model, client=client) @@ -54,7 +52,7 @@ def __repr__(self) -> str: class Plan: """Callable client-side reference to a registered BlueAPI plan.""" - def __init__(self, name: str, model: PlanModel, client: ClientProtocol): + def __init__(self, name, model: PlanModel, client: ClientProtocol): self.name = name self.model = model self._client = client @@ -63,9 +61,6 @@ def __init__(self, name: str, model: PlanModel, client: ClientProtocol): def __call__(self, *args, **kwargs) -> Any: """Execute the plan with the supplied arguments. - Positional arguments are mapped to parameters in the order defined by - the plan's parameter schema. Keyword arguments are passed by name. - Returns: The result returned by the plan. @@ -74,7 +69,7 @@ def __call__(self, *args, **kwargs) -> Any: TypeError: If the supplied arguments do not match the plan parameter schema. """ - req = TaskRequest( + req = TaskRequestV2( name=self.name, params=self._build_args(*args, **kwargs), instrument_session=self._client.instrument_session, @@ -97,12 +92,11 @@ def properties(self) -> dict[str, Any]: def required(self) -> list[str]: return self.model.parameter_schema.get("required", []) - def _build_args(self, *args, **kwargs) -> dict[str, Any]: - """Build a parameter mapping from positional and keyword arguments. + def _build_args(self, *args: Any, **kwargs: Any) -> TaskParametersV2: + """Check parameters against schema and build TaskParametersV2. - Positional arguments are assigned to parameters according to their - order in the plan's parameter schema. Keyword arguments are then - added by name. + Returns: + TaskParametersV2 Raises: TypeError: If too many positional arguments, unexpected keyword @@ -115,28 +109,26 @@ def _build_args(self, *args, **kwargs) -> dict[str, Any]: args, kwargs, ) + properties = list(self.properties) - if len(args) > len(self.properties): + if len(args) > len(properties): raise TypeError(f"{self.name} got too many arguments") - if extra := {k for k in kwargs if k not in self.properties}: + + if extra := {k for k in kwargs if k not in properties}: raise TypeError(f"{self.name} got unexpected arguments: {extra}") - params = {} - # Initially fill parameters using positional args assuming the order - # from the parameter_schema - for req, arg in zip(self.properties, args, strict=False): - params[req] = arg + positional_names = properties[: len(args)] + + if duplicate := set(positional_names) & kwargs.keys(): + name = next(iter(duplicate)) + raise TypeError(f"{self.name} got multiple values for {name}") - # Then append any values given via kwargs - for key, value in kwargs.items(): - # If we've already assumed a positional arg was this value, bail out - if key in params: - raise TypeError(f"{self.name} got multiple values for {key}") - params[key] = value + supplied = set(positional_names) | kwargs.keys() - if missing := {k for k in self.required if k not in params}: + if missing := set(self.required) - supplied: raise TypeError(f"Missing argument(s) for {missing}") - return params + + return TaskParametersV2(args=args, kwargs=kwargs) def __repr__(self) -> str: """Return a signature-like representation of the plan. diff --git a/src/blueapi/client/protocols.py b/src/blueapi/client/protocols.py index dbb923af6..235d7b07f 100644 --- a/src/blueapi/client/protocols.py +++ b/src/blueapi/client/protocols.py @@ -1,6 +1,6 @@ from typing import Protocol -from blueapi.service.model import DeviceModel, DeviceResponse, TaskRequest +from blueapi.service.model import DeviceModel, DeviceResponse, TaskRequestV2 from blueapi.worker.event import TaskStatus from .event_bus import OnAnyEvent @@ -12,7 +12,7 @@ class ClientProtocol(Protocol): def run_task( self, - task: TaskRequest, + task: TaskRequestV2, on_event: OnAnyEvent | None = None, timeout: float | None = None, ) -> TaskStatus: ... diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py index da6b4b6dd..42b201e31 100644 --- a/src/blueapi/client/rest.py +++ b/src/blueapi/client/rest.py @@ -30,6 +30,8 @@ PythonEnvironmentResponse, SourceInfo, TaskRequest, + TaskRequestV1, + TaskRequestV2, TaskResponse, TasksListResponse, WorkerTask, @@ -253,6 +255,8 @@ def get_active_task(self) -> WorkerTask: return self._request_and_deserialize("/worker/task", WorkerTask) def create_task(self, task: TaskRequest) -> TaskResponse: + if isinstance(task, TaskRequestV1): + task = TaskRequestV2.from_v1(task) return self._request_and_deserialize( "/tasks", TaskResponse, @@ -359,6 +363,9 @@ def _request_and_deserialize( def run_blocking( self, req: TaskRequest ) -> Iterable[DataEvent | WorkerEvent | ProgressEvent]: + if isinstance(req, TaskRequestV1): + req = TaskRequestV2.from_v1(req) + url = self._config.ws_address.unicode_string().rstrip("/") + "/api/v2/run_plan" headers = get_context_propagator() if self.session_manager: diff --git a/src/blueapi/config.py b/src/blueapi/config.py index 9c339c71f..f65dd65f0 100644 --- a/src/blueapi/config.py +++ b/src/blueapi/config.py @@ -335,7 +335,7 @@ class ApplicationConfig(BlueapiBaseModel): """ #: API version to publish in OpenAPI schema - REST_API_VERSION: ClassVar[str] = "1.5.0" + REST_API_VERSION: ClassVar[str] = "1.6.0" LICENSE_INFO: ClassVar[dict[str, str]] = { "name": "Apache 2.0", diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index 76ce74e64..c59a70b67 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -439,19 +439,30 @@ def _type_spec_for_function( Returns: Mapping of {name: (type, default)} to be used by pydantic for deserialising - function arguments + function arguments """ args = signature(func).parameters types = get_type_hints(func) new_args: dict[str, tuple[type, FieldInfo]] = {} + for name, para in args.items(): arg_type = types.get(name, Parameter.empty) + if arg_type is Parameter.empty: raise ValueError( f"Type annotation is required for '{name}' in '{func.__name__}'" ) - no_default = para.default is Parameter.empty + match para.kind: + case Parameter.VAR_POSITIONAL: + arg_type = tuple[arg_type, ...] + case Parameter.VAR_KEYWORD: + arg_type = dict[str, arg_type] + + no_default = para.default is Parameter.empty and para.kind not in ( + Parameter.VAR_POSITIONAL, + Parameter.VAR_KEYWORD, + ) if ( isclass(arg_type) and (issubclass(arg_type, BaseModel) or is_dataclass(arg_type)) @@ -462,9 +473,15 @@ def _type_spec_for_function( info = FieldInfo(default_factory=default_factory) else: _type = self._convert_type(arg_type, no_default) + match para.default: case Parameter.empty: - info = FieldInfo(default_factory=None) + if para.kind is Parameter.VAR_POSITIONAL: + info = FieldInfo(default_factory=tuple) + elif para.kind is Parameter.VAR_KEYWORD: + info = FieldInfo(default_factory=dict) + else: + info = FieldInfo(default_factory=None) case None: info = FieldInfo(default_factory=lambda: None) case _: diff --git a/src/blueapi/service/interface.py b/src/blueapi/service/interface.py index 6c0c5befb..bc9a58184 100644 --- a/src/blueapi/service/interface.py +++ b/src/blueapi/service/interface.py @@ -23,6 +23,8 @@ PythonEnvironmentResponse, SourceInfo, TaskRequest, + TaskRequestV1, + TaskRequestV2, WorkerTask, ) from blueapi.utils.serialization import access_blob @@ -156,6 +158,9 @@ def submit_task( task_request: TaskRequest, metadata: dict[str, Any] | None = None ) -> str: """Submit a task to be run on begin_task""" + if isinstance(task_request, TaskRequestV1): + task_request = TaskRequestV2.from_v1(task_request) + # Can't default arg to mutable data structure: if metadata is None: metadata = {} diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py index 0d8fbf477..64b01cbd5 100644 --- a/src/blueapi/service/main.py +++ b/src/blueapi/service/main.py @@ -55,7 +55,7 @@ Unauthorized, Update, ) -from blueapi.worker import TrackableTask, WorkerState +from blueapi.worker import TaskParametersV2, TrackableTask, WorkerState from blueapi.worker.event import ProgressEvent, TaskStatusEnum, WorkerEvent from blueapi.worker.worker_errors import WorkerBusyError @@ -78,6 +78,7 @@ SourceInfo, StateChangeRequest, TaskRequest, + TaskRequestV2, TaskResponse, TasksListResponse, WorkerTask, @@ -318,9 +319,9 @@ def get_device_by_name( return runner.run(interface.get_device, name) -example_task_request = TaskRequest( +example_task_request = TaskRequestV2( name="count", - params={"detectors": ["x"]}, + params=TaskParametersV2(kwargs={"detectors": ["x"]}), instrument_session="cm12345-1", ) diff --git a/src/blueapi/service/model.py b/src/blueapi/service/model.py index dbaa1d965..b29b943ad 100644 --- a/src/blueapi/service/model.py +++ b/src/blueapi/service/model.py @@ -1,5 +1,5 @@ import uuid -from collections.abc import Iterable, Mapping +from collections.abc import Iterable from enum import StrEnum from typing import Annotated, Any @@ -11,7 +11,7 @@ from blueapi.core import BLUESKY_PROTOCOLS, Device, Plan from blueapi.core.context import generic_bounds from blueapi.utils import BlueapiBaseModel -from blueapi.worker import WorkerState +from blueapi.worker import TaskParametersV1, TaskParametersV2, WorkerState from blueapi.worker.task_worker import TaskWorker, TrackableTask _UNKNOWN_NAME = "UNKNOWN" @@ -58,13 +58,13 @@ class TasksListResponse(BlueapiBaseModel): tasks: list[TrackableTask] = Field(description="List of tasks") -class TaskRequest(BlueapiBaseModel): +class TaskRequestV1(BlueapiBaseModel): """ - Request to run a task with related info + Request to run a task with related info. Supports **kwargs only. """ name: str = Field(description="Name of plan to run") - params: Mapping[str, Any] = Field( + params: TaskParametersV1 = Field( description="Values for parameters to plan, if any", default_factory=dict ) instrument_session: str = Field( @@ -72,6 +72,32 @@ class TaskRequest(BlueapiBaseModel): ) +class TaskRequestV2(BlueapiBaseModel): + """ + Request to run a task with related info. Supports *args and **kwargs. + """ + + name: str = Field(description="Name of plan to run") + params: TaskParametersV2 = Field( + description="Values for parameters to plan, if any", + default_factory=TaskParametersV2, + ) + instrument_session: str = Field( + description="Instrument session associated with this task", + ) + + @classmethod + def from_v1(cls, task_v1: TaskRequestV1) -> "TaskRequestV2": + return cls( + name=task_v1.name, + instrument_session=task_v1.instrument_session, + params=TaskParametersV2(kwargs=task_v1.params), + ) + + +TaskRequest = TaskRequestV2 | TaskRequestV1 + + class DeviceRequest(BlueapiBaseModel): """ A query for devices diff --git a/src/blueapi/worker/__init__.py b/src/blueapi/worker/__init__.py index 85ae49b45..6c0aaacb4 100644 --- a/src/blueapi/worker/__init__.py +++ b/src/blueapi/worker/__init__.py @@ -1,11 +1,13 @@ from .event import ProgressEvent, StatusView, TaskStatus, WorkerEvent, WorkerState -from .task import Task +from .task import Task, TaskParametersV1, TaskParametersV2 from .task_worker import TaskWorker, TrackableTask from .worker_errors import WorkerAlreadyStartedError, WorkerBusyError __all__ = [ "TaskWorker", "Task", + "TaskParametersV1", + "TaskParametersV2", "WorkerEvent", "WorkerState", "StatusView", diff --git a/src/blueapi/worker/task.py b/src/blueapi/worker/task.py index e2efade0d..467713658 100644 --- a/src/blueapi/worker/task.py +++ b/src/blueapi/worker/task.py @@ -1,8 +1,9 @@ import logging -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from inspect import Parameter, signature from typing import Any -from pydantic import BaseModel, Field, TypeAdapter +from pydantic import Field, TypeAdapter from blueapi.core import BlueskyContext from blueapi.utils import BlueapiBaseModel @@ -10,55 +11,101 @@ LOGGER = logging.getLogger(__name__) +TaskParametersV1 = Mapping[str, Any] + + +class TaskParametersV2(BlueapiBaseModel): + args: Sequence[Any] = Field(default_factory=list) + kwargs: Mapping[str, Any] = Field(default_factory=dict) + + @classmethod + def from_v1(cls, param: TaskParametersV1): + return cls(kwargs=param) + + class Task(BlueapiBaseModel): """ Task that will run a plan """ name: str = Field(description="Name of plan to run") - params: Mapping[str, Any] = Field( - description="Values for parameters to plan, if any", default_factory=dict + params: TaskParametersV2 = Field( + description="Values for parameters to plan, if any", + default_factory=TaskParametersV2, ) metadata: dict[str, Any] = Field( description="Any metadata to apply to all runs within this task", default_factory=dict, ) - def prepare_params(self, ctx: BlueskyContext) -> Mapping[str, Any]: - model = _lookup_params(ctx, self) - # Re-create dict manually to avoid nesting in model_dump output - return {field: getattr(model, field) for field in model.__pydantic_fields__} + def prepare_params( + self, ctx: BlueskyContext + ) -> tuple[list[Any], Mapping[str, Any]]: + """ + Checks the configured plan parameters against context + + Args: + ctx: Context holding plans and devices + + Returns: + tuple[list[Any], Mapping[str, Any]]: The prepared *args and **kwargs + parameters for the plan. + """ + plan = ctx.plans[self.name] + func = ctx.plan_functions[self.name] + + sig = signature(func) + bound = sig.bind_partial(*self.params.args, **self.params.kwargs) + + # Only validate explicitly supplied arguments. This allows Pydantic's + # default/default_factory to provide injected defaults. + adapter = TypeAdapter(plan.model) + validated = adapter.validate_python(bound.arguments) + + args: list[Any] = [] + kwargs: dict[str, Any] = {} + + for name, parameter in sig.parameters.items(): + supplied = name in bound.arguments + + if supplied: + value = getattr(validated, name) + + match parameter.kind: + case Parameter.POSITIONAL_ONLY: + args.append(value) + + case Parameter.POSITIONAL_OR_KEYWORD: + if name in self.params.kwargs: + kwargs[name] = value + else: + args.append(value) + + case Parameter.VAR_POSITIONAL: + args.extend(value) + + case Parameter.KEYWORD_ONLY: + kwargs[name] = value + + case Parameter.VAR_KEYWORD: + kwargs.update(value) + else: + # Let the generated Pydantic model provide the default. + value = getattr(validated, name) + kwargs[name] = value + + return args, kwargs def do_task(self, ctx: BlueskyContext) -> None: LOGGER.info( f"Asked to run plan {self.name} with {self.params} and " f"metadata {self.metadata} for all runs" ) - - func = ctx.plan_functions[self.name] - prepared_params = self.prepare_params(ctx) + plan = ctx.plan_functions[self.name] + prepared_args, prepared_kwargs = self.prepare_params(ctx) ctx.run_engine.md.update(self.metadata) - result = ctx.run_engine(func(**prepared_params)) + result = ctx.run_engine(plan(*prepared_args, **prepared_kwargs)) if isinstance(result, tuple): # pragma: no cover # this is never true if the run_engine is configured correctly return None return result.plan_result - - -def _lookup_params(ctx: BlueskyContext, task: Task) -> BaseModel: - """ - Checks plan parameters against context - - Args: - ctx: Context holding plans and devices - plan: Plan object including schema - params: Parameter values to be validated against schema - - Returns: - Mapping[str, Any]: _description_ - """ - - plan = ctx.plans[task.name] - model = plan.model - adapter = TypeAdapter(model) - return adapter.validate_python(task.params) diff --git a/tests/system_tests/test_blueapi_system.py b/tests/system_tests/test_blueapi_system.py index 8c4dc19cf..6195173f3 100644 --- a/tests/system_tests/test_blueapi_system.py +++ b/tests/system_tests/test_blueapi_system.py @@ -32,7 +32,9 @@ from blueapi.service.model import ( DeviceResponse, PlanResponse, - TaskRequest, + TaskParametersV2, + TaskRequestV1, + TaskRequestV2, TaskResponse, WorkerTask, ) @@ -119,10 +121,10 @@ def scan_id(local_numtracker) -> int: def task_factory( user: ValidUser, instrument_session: str | None, time: float = 0.0 -) -> TaskRequest: - return TaskRequest( +) -> TaskRequestV2: + return TaskRequestV2( name="sleep", - params={"time": time}, + params=TaskParametersV2(kwargs={"time": time}), instrument_session=instrument_session if instrument_session else VALID_INSTRUMENT_SESSION[user], @@ -130,12 +132,12 @@ def task_factory( @pytest.fixture -def small_task(user: ValidUser, instrument_session: str | None) -> TaskRequest: +def small_task(user: ValidUser, instrument_session: str | None) -> TaskRequestV2: return task_factory(user, instrument_session) @pytest.fixture -def long_task(user: ValidUser, instrument_session: str | None) -> TaskRequest: +def long_task(user: ValidUser, instrument_session: str | None) -> TaskRequestV2: return task_factory(user, instrument_session, time=1.0) @@ -339,14 +341,14 @@ def test_client_non_existent_device(client: BlueapiClient): def test_create_task_and_delete_task_by_id( - rest_client: BlueapiRestClient, small_task: TaskRequest + rest_client: BlueapiRestClient, small_task: TaskRequestV2 ): create_task = rest_client.create_task(small_task) rest_client.clear_task(create_task.task_id) def test_task_metadata_propagated( - rest_client: BlueapiRestClient, small_task: TaskRequest + rest_client: BlueapiRestClient, small_task: TaskRequestV2 ): response = rest_client.create_task(small_task) trackable_task = rest_client.get_task(response.task_id) @@ -363,16 +365,16 @@ def test_task_metadata_propagated( def test_create_task_validation_error(rest_client: BlueapiRestClient): with pytest.raises(BlueskyRequestError, match="Internal Server Error"): rest_client.create_task( - TaskRequest( + TaskRequestV2( name="Not-exists", - params={"Not-exists": 0.0}, + params=TaskParametersV2(kwargs={"Not-exists": 0.0}), instrument_session="Not-exists", ) ) def test_get_all_tasks( - rest_client: BlueapiRestClient, small_task: TaskRequest, long_task: TaskRequest + rest_client: BlueapiRestClient, small_task: TaskRequestV2, long_task: TaskRequestV2 ): created_tasks: list[TaskResponse] = [] for task in [small_task, long_task]: @@ -389,7 +391,7 @@ def test_get_all_tasks( rest_client.clear_task(task_id) -def test_get_task_by_id(rest_client: BlueapiRestClient, small_task: TaskRequest): +def test_get_task_by_id(rest_client: BlueapiRestClient, small_task: TaskRequestV2): created_task = rest_client.create_task(small_task) get_task = rest_client.get_task(created_task.task_id) @@ -413,7 +415,7 @@ def test_delete_non_existent_task(rest_client: BlueapiRestClient): rest_client.clear_task("Not-exists") -def test_put_worker_task(rest_client: BlueapiRestClient, small_task: TaskRequest): +def test_put_worker_task(rest_client: BlueapiRestClient, small_task: TaskRequestV2): created_task = rest_client.create_task(small_task) rest_client.update_worker_task(WorkerTask(task_id=created_task.task_id)) active_task = rest_client.get_active_task() @@ -422,7 +424,7 @@ def test_put_worker_task(rest_client: BlueapiRestClient, small_task: TaskRequest def test_put_worker_task_fails_if_not_idle( - rest_client: BlueapiRestClient, small_task: TaskRequest, long_task: TaskRequest + rest_client: BlueapiRestClient, small_task: TaskRequestV2, long_task: TaskRequestV2 ): small_task_id = rest_client.create_task(small_task).task_id long_task_id = rest_client.create_task(long_task).task_id @@ -452,7 +454,7 @@ def test_set_state_transition_error(client: BlueapiClient): assert "Cannot transition from IDLE to PAUSED" in exception.value.args[1] -def test_get_task_by_status(rest_client: BlueapiRestClient, small_task: TaskRequest): +def test_get_task_by_status(rest_client: BlueapiRestClient, small_task: TaskRequestV2): task_1 = rest_client.create_task(small_task) task_2 = rest_client.create_task(small_task) task_by_pending = rest_client.get_all_tasks() @@ -483,7 +485,9 @@ def test_get_task_by_status(rest_client: BlueapiRestClient, small_task: TaskRequ rest_client.clear_task(task_id=task_2.task_id) -def test_progress_with_stomp(client_with_stomp: BlueapiClient, small_task: TaskRequest): +def test_progress_with_stomp( + client_with_stomp: BlueapiClient, small_task: TaskRequestV2 +): all_events: list[AnyEvent] = [] def on_event(event: AnyEvent): @@ -537,18 +541,18 @@ def test_delete_current_environment(client: BlueapiClient): @pytest.mark.parametrize( - "task,user", + "task", [ - ( - TaskRequest( + pytest.param( + TaskRequestV1( name="count", params={"detectors": ["det"], "num": 3}, - instrument_session=VALID_INSTRUMENT_SESSION[User.alice], + instrument_session="", # Set below ), - User.alice, + id="v1", ), - ( - TaskRequest( + pytest.param( + TaskRequestV1( name="spec_scan", params={ "detectors": ["det"], @@ -571,20 +575,109 @@ def test_delete_current_environment(client: BlueapiClient): "type": "Product", }, }, - instrument_session=VALID_INSTRUMENT_SESSION[User.bob], + instrument_session="", + ), + id="spec_scan-v1", + ), + pytest.param( + TaskRequestV2( + name="count", + params=TaskParametersV2(args=[["det"], 3]), + instrument_session="", + ), + id="v2_args", + ), + pytest.param( + TaskRequestV2( + name="count", + params=TaskParametersV2(kwargs={"detectors": ["det"], "num": 3}), + instrument_session="", + ), + id="v2_kwargs", + ), + pytest.param( + TaskRequestV2( + name="count", + params=TaskParametersV2(args=[["det"]], kwargs={"num": 3}), + instrument_session="", + ), + id="v2_args_kwargs", + ), + pytest.param( + TaskRequestV2( + name="spec_scan", + params=TaskParametersV2( + args=[["det"]], + kwargs={ + "spec": { + "outer": { + "axis": "stage.x", + "start": 0.0, + "stop": 0.4, + "num": 2, + "type": "Linspace", + }, + "inner": { + "axis": "stage.theta", + "start": 5.0, + "stop": 5.3, + "num": 3, + "type": "Linspace", + }, + "gap": True, + "type": "Product", + }, + }, + ), + instrument_session="", ), - User.bob, + id="spec_scan-v2_args_kwargs", + ), + pytest.param( + TaskRequestV2( + name="spec_scan", + params=TaskParametersV2( + kwargs={ + "detectors": ["det"], + "spec": { + "outer": { + "axis": "stage.x", + "start": 0.0, + "stop": 0.4, + "num": 2, + "type": "Linspace", + }, + "inner": { + "axis": "stage.theta", + "start": 5.0, + "stop": 5.3, + "num": 3, + "type": "Linspace", + }, + "gap": True, + "type": "Product", + }, + } + ), + instrument_session="", + ), + id="spec_scan-v2_kwargs", ), ], ) +@pytest.mark.parametrize("user", [User.alice, User.bob]) @pytest.mark.parametrize("run_method", ["run_task", "run_blocking"]) def test_plan_runs( client_with_stomp: BlueapiClient, - task: TaskRequest, + task: TaskRequestV2, scan_id: int, user: ValidUser, run_method: str, ): + # Set instrument session after creation so we can test every single task combination + # with every user + task.instrument_session = VALID_INSTRUMENT_SESSION[user] + resource = Queue(maxsize=1) start = Queue(maxsize=1) @@ -636,17 +729,41 @@ def on_event(event: AnyEvent) -> None: @pytest.mark.parametrize( "task", [ - TaskRequest( - name="set_absolute", - params={ - "movable": "stage.x", - "value": 1.0, - }, - instrument_session=VALID_INSTRUMENT_SESSION[User.alice], + pytest.param( + TaskRequestV1( + name="set_absolute", + params={"movable": "stage.x", "value": 1.0}, + instrument_session=VALID_INSTRUMENT_SESSION[User.alice], + ), + id="v1", + ), + pytest.param( + TaskRequestV2( + name="set_absolute", + params=TaskParametersV2(args=["stage.x", 1.0]), + instrument_session=VALID_INSTRUMENT_SESSION[User.alice], + ), + id="v2_args", + ), + pytest.param( + TaskRequestV2( + name="set_absolute", + params=TaskParametersV2(args=["stage.x"], kwargs={"value": 1.0}), + instrument_session=VALID_INSTRUMENT_SESSION[User.alice], + ), + id="v2_args_kwargs", + ), + pytest.param( + TaskRequestV2( + name="set_absolute", + params=TaskParametersV2(kwargs={"movable": "stage.x", "value": 1.0}), + instrument_session=VALID_INSTRUMENT_SESSION[User.alice], + ), + id="v2_kwargs", ), ], ) -def test_stub_runs(client_with_stomp: BlueapiClient, task: TaskRequest): +def test_stub_runs(client_with_stomp: BlueapiClient, task: TaskRequestV2): final_event = client_with_stomp.run_task(task) assert isinstance(final_event.result, TaskResult) assert final_event.task_complete @@ -661,13 +778,13 @@ def test_task_submission_after_invalid_task(client_with_stomp: BlueapiClient): # ...but should leave the serve in a state where it can still run tasks res = client_with_stomp.run_task( - TaskRequest( + TaskRequestV2( name="count", - params={ - "detectors": [ - "det", - ], - }, + params=TaskParametersV2( + kwargs={ + "detectors": ["det"], + } + ), instrument_session=VALID_INSTRUMENT_SESSION[User.alice], ) ) @@ -688,7 +805,7 @@ def test_task_submission_after_invalid_task(client_with_stomp: BlueapiClient): @pytest.mark.parametrize("method_name", ["create_task", "run_blocking"]) def test_run_task_without_authz( client: BlueapiClient, - small_task: TaskRequest, + small_task: TaskRequestV2, user: ValidUser, instrument_session: str, method_name: str, @@ -713,7 +830,7 @@ def test_run_task_without_authz( ) def test_create_task_authorization( client: BlueapiClient, - small_task: TaskRequest, + small_task: TaskRequestV2, user: ValidUser, instrument_session: str, ): @@ -881,7 +998,7 @@ def test_admin_can_abort_any_task( def test_run_blocking_requires_auth( - client_without_auth: BlueapiClient, small_task: TaskRequest + client_without_auth: BlueapiClient, small_task: TaskRequestV2 ): with pytest.raises(UnauthorisedAccessError): client_without_auth.run_blocking(small_task) diff --git a/tests/unit_tests/cli/test_cli.py b/tests/unit_tests/cli/test_cli.py index 28cb4fbc9..aa5f2c2c6 100644 --- a/tests/unit_tests/cli/test_cli.py +++ b/tests/unit_tests/cli/test_cli.py @@ -47,9 +47,10 @@ PlanModel, PlanResponse, PythonEnvironmentResponse, - TaskRequest, + TaskRequestV2, TaskResponse, ) +from blueapi.worker import TaskParametersV2 from blueapi.worker.event import ( ProgressEvent, TaskError, @@ -210,7 +211,7 @@ def test_invalid_config_via_env(runner: CliRunner): def test_submit_plan(runner: CliRunner): body_data = { "name": "sleep", - "params": {"time": 5}, + "params": {"args": [], "kwargs": {"time": 5}}, "instrument_session": "cm12345-1", } @@ -260,9 +261,18 @@ def test_submit_plan_without_stomp(runner: CliRunner): ) +@pytest.mark.parametrize( + "params, expected", + [ + [{"time": 3}, {"args": [], "kwargs": {"time": 3}}], + [{"args": [], "kwargs": {"time": 3}}, {"args": [], "kwargs": {"time": 3}}], + [{"args": [3], "kwargs": {}}, {"args": [3], "kwargs": {}}], + ], + ids=["v1", "v2_kwargs", "v2_args"], +) @patch("blueapi.client.client.StompClient") @responses.activate -def test_run_plan(stomp_client: StompClient, runner: CliRunner): +def test_run_plan(stomp_client: StompClient, runner: CliRunner, params, expected): task_id = "abcd-1234" submit_response = responses.post( url="http://a.fake.host:12345/tasks", @@ -270,7 +280,7 @@ def test_run_plan(stomp_client: StompClient, runner: CliRunner): matchers.json_params_matcher( { "name": "sleep", - "params": {"time": 3}, + "params": expected, "instrument_session": "cm12345-1", } ) @@ -339,11 +349,11 @@ def mock_events(topic: MessageTopic, callback: Callable[[Any, Any], Any]): "-i", "cm12345-1", "sleep", - '{"time": 3}', + json.dumps(params), ], ) - assert result.exit_code == 0 + assert result.exit_code == 0, result.output assert submit_response.call_count == 1 assert run_response.call_count == 1 @@ -386,7 +396,9 @@ def test_run_plan_feedback( ) bc.add_callback.assert_called_once() bc.run_task.assert_called_once_with( - TaskRequest(name="name", params={}, instrument_session="cm12345-1"), + TaskRequestV2( + name="name", params=TaskParametersV2(), instrument_session="cm12345-1" + ), ) assert res.exit_code == 0 assert res.stdout == message @@ -400,7 +412,7 @@ def test_run_plan_background_without_stomp(runner: CliRunner): matchers.json_params_matcher( { "name": "sleep", - "params": {"time": 3}, + "params": {"args": [], "kwargs": {"time": 3}}, "instrument_session": "cm12345-1", } ) @@ -501,9 +513,9 @@ def test_can_pass_an_instrument_session_with_an_environment_variable( ) assert result.exit_code == 0, result.output mock_create_task.assert_called_once_with( - TaskRequest( + TaskRequestV2( name="sleep", - params={"time": 5.0}, + params=TaskParametersV2(kwargs={"time": 5.0}), instrument_session="cm12345-1", ) ) @@ -1420,10 +1432,27 @@ def test_config_schema( assert json.loads(result.output) == expected -@pytest.mark.parametrize("value,result", [({}, {}), ("{}", {}), (None, None)]) -def test_task_parameter_type(value, result): +@pytest.mark.parametrize( + "value, expected", + [ + ({}, TaskParametersV2()), + ("{}", TaskParametersV2()), + (None, TaskParametersV2()), + ('{"time": 3}', TaskParametersV2(kwargs={"time": 3})), + ( + '{"args": [], "kwargs": {"time": 3}}', + TaskParametersV2(args=[], kwargs={"time": 3}), + ), + ( + '{"args": [3], "kwargs": {}}', + TaskParametersV2(args=[3], kwargs={}), + ), + ], + ids=["empty-dict", "empty-json", "none", "v1", "v2_kwargs", "v2_args"], +) +def test_task_parameter_type(value, expected): t = ParametersType() - assert t.convert(value, None, None) == result + assert t.convert(value, None, None) == expected @pytest.mark.parametrize( @@ -1489,6 +1518,8 @@ def test_run_ws_runs_blocking_plan(mock_client: Mock, runner: CliRunner): bc.add_callback.assert_called_once() bc.run_task.assert_not_called() bc.run_blocking.assert_called_once_with( - TaskRequest(name="name", params={}, instrument_session="cm12345-1"), + TaskRequestV2( + name="name", params=TaskParametersV2(), instrument_session="cm12345-1" + ), ) assert res.exit_code == 0 diff --git a/tests/unit_tests/client/constants.py b/tests/unit_tests/client/constants.py index fad6057c6..727cd15dc 100644 --- a/tests/unit_tests/client/constants.py +++ b/tests/unit_tests/client/constants.py @@ -11,6 +11,7 @@ ) from blueapi.worker import Task, TrackableTask, WorkerEvent, WorkerState from blueapi.worker.event import TaskError, TaskResult, TaskStatus +from blueapi.worker.task import TaskParametersV2 PLANS = PlanResponse( plans=[ @@ -39,7 +40,7 @@ ] ) DEVICE = DeviceModel(name="foo", protocols=[]) -TASK = TrackableTask(task_id="foo", task=Task(name="bar", params={})) +TASK = TrackableTask(task_id="foo", task=Task(name="bar", params=TaskParametersV2())) TASKS = TasksListResponse(tasks=[TASK]) ACTIVE_TASK = WorkerTask(task_id="bar") ENVIRONMENT_ID = uuid.uuid4() diff --git a/tests/unit_tests/client/test_client.py b/tests/unit_tests/client/test_client.py index 3293c3610..655449efc 100644 --- a/tests/unit_tests/client/test_client.py +++ b/tests/unit_tests/client/test_client.py @@ -36,7 +36,7 @@ PlanModel, PlanResponse, ProtocolInfo, - TaskRequest, + TaskRequestV2, TaskResponse, WorkerTask, ) @@ -145,10 +145,10 @@ def test_create_and_start_task_calls_both_creating_and_starting_endpoints( mock_rest.create_task.return_value = TaskResponse(task_id="baz") mock_rest.update_worker_task.return_value = TaskResponse(task_id="baz") client.create_and_start_task( - TaskRequest(name="baz", instrument_session="cm12345-1") + TaskRequestV2(name="baz", instrument_session="cm12345-1") ) mock_rest.create_task.assert_called_once_with( - TaskRequest(name="baz", instrument_session="cm12345-1") + TaskRequestV2(name="baz", instrument_session="cm12345-1") ) mock_rest.update_worker_task.assert_called_once_with(WorkerTask(task_id="baz")) @@ -160,7 +160,7 @@ def test_create_and_start_task_fails_if_task_creation_fails( mock_rest.create_task.side_effect = BlueskyRemoteControlError("No can do") with pytest.raises(BlueskyRemoteControlError): client.create_and_start_task( - TaskRequest(name="baz", instrument_session="cm12345-1") + TaskRequestV2(name="baz", instrument_session="cm12345-1") ) @@ -172,7 +172,7 @@ def test_create_and_start_task_fails_if_task_id_is_wrong( mock_rest.update_worker_task.return_value = TaskResponse(task_id="bar") with pytest.raises(BlueskyRemoteControlError): client.create_and_start_task( - TaskRequest(name="baz", instrument_session="cm12345-1") + TaskRequestV2(name="baz", instrument_session="cm12345-1") ) @@ -184,7 +184,7 @@ def test_create_and_start_task_fails_if_task_start_fails( mock_rest.update_worker_task.side_effect = BlueskyRemoteControlError("No can do") with pytest.raises(BlueskyRemoteControlError): client.create_and_start_task( - TaskRequest(name="baz", instrument_session="cm12345-1") + TaskRequestV2(name="baz", instrument_session="cm12345-1") ) @@ -353,7 +353,7 @@ def test_cannot_run_task_without_message_bus(client: BlueapiClient): MissingStompConfigurationError, match="Stomp configuration required to run plans is missing or disabled", ): - client.run_task(TaskRequest(name="foo", instrument_session="cm12345-1")) + client.run_task(TaskRequestV2(name="foo", instrument_session="cm12345-1")) def test_run_task_sets_up_control( @@ -367,9 +367,11 @@ def test_run_task_sets_up_control( ctx.correlation_id = "foo" mock_events.subscribe_to_all_events = lambda on_event: on_event(COMPLETE_EVENT, ctx) - client_with_events.run_task(TaskRequest(name="foo", instrument_session="cm12345-1")) + client_with_events.run_task( + TaskRequestV2(name="foo", instrument_session="cm12345-1") + ) mock_rest.create_task.assert_called_once_with( - TaskRequest(name="foo", instrument_session="cm12345-1") + TaskRequestV2(name="foo", instrument_session="cm12345-1") ) mock_rest.update_worker_task.assert_called_once_with(WorkerTask(task_id="foo")) @@ -388,7 +390,7 @@ def test_run_task_fails_on_failing_event( on_event = Mock() outcome = client_with_events.run_task( - TaskRequest(name="foo", instrument_session="cm12345-1"), + TaskRequestV2(name="foo", instrument_session="cm12345-1"), on_event=on_event, ) assert outcome.task_failed @@ -436,7 +438,8 @@ def callback(on_event: Callable[[AnyEvent, MessageContext], None]): mock_on_event = Mock() client_with_events.run_task( - TaskRequest(name="foo", instrument_session="cm12345-1"), on_event=mock_on_event + TaskRequestV2(name="foo", instrument_session="cm12345-1"), + on_event=mock_on_event, ) assert mock_on_event.mock_calls == [call(test_event), call(COMPLETE_EVENT)] @@ -478,7 +481,8 @@ def callback(on_event: Callable[[AnyEvent, MessageContext], None]): mock_on_event = Mock() client_with_events.run_task( - TaskRequest(name="foo", instrument_session="cm12345-1"), on_event=mock_on_event + TaskRequestV2(name="foo", instrument_session="cm12345-1"), + on_event=mock_on_event, ) mock_on_event.assert_called_once_with(COMPLETE_EVENT) @@ -566,7 +570,7 @@ def test_create_and_start_task_span_ok( mock_rest.update_worker_task.return_value = TaskResponse(task_id="baz") with asserting_span_exporter(exporter, "create_and_start_task", "task"): client.create_and_start_task( - TaskRequest(name="baz", instrument_session="cm12345-1") + TaskRequestV2(name="baz", instrument_session="cm12345-1") ) @@ -631,7 +635,7 @@ def test_cannot_run_task_span_ok( match="Stomp configuration required to run plans is missing or disabled", ): with asserting_span_exporter(exporter, "grun_task"): - client.run_task(TaskRequest(name="foo", instrument_session="cm12345-1")) + client.run_task(TaskRequestV2(name="foo", instrument_session="cm12345-1")) def test_instrument_session_required(client): @@ -698,7 +702,9 @@ def subscribe(on_event: Callable[[AnyEvent, MessageContext], None]): mock_events.subscribe_to_all_events = subscribe # type: ignore - client_with_events.run_task(TaskRequest(name="foo", instrument_session="cm12345-1")) + client_with_events.run_task( + TaskRequestV2(name="foo", instrument_session="cm12345-1") + ) assert callback.mock_calls == [call(test_event), call(COMPLETE_EVENT)] @@ -726,7 +732,9 @@ def subscribe(on_event: Callable[[AnyEvent, MessageContext], None]): mock_events.subscribe_to_all_events = subscribe # type: ignore - client_with_events.run_task(TaskRequest(name="foo", instrument_session="cm12345-1")) + client_with_events.run_task( + TaskRequestV2(name="foo", instrument_session="cm12345-1") + ) assert failing_callback.mock_calls == [call(evt), call(COMPLETE_EVENT)] assert callback.mock_calls == [call(evt), call(COMPLETE_EVENT)] @@ -780,7 +788,7 @@ def test_client_logout(mock_rest: Mock, client: BlueapiClient): @pytest.mark.parametrize("event", [COMPLETE_EVENT, FAILED_EVENT]) def test_run_blocking(event: WorkerEvent, client: BlueapiClient, mock_rest: Mock): mock_rest.run_blocking.side_effect = lambda req: [event] - res = client.run_blocking(TaskRequest(name="foo", instrument_session="cm12345-1")) + res = client.run_blocking(TaskRequestV2(name="foo", instrument_session="cm12345-1")) assert res == event.task_status diff --git a/tests/unit_tests/client/test_plan_cache.py b/tests/unit_tests/client/test_plan_cache.py index 3da4d1c9b..81fca4146 100644 --- a/tests/unit_tests/client/test_plan_cache.py +++ b/tests/unit_tests/client/test_plan_cache.py @@ -5,7 +5,7 @@ from blueapi.client.client import PlanCache from blueapi.client.plan_cache import Plan -from blueapi.service.model import PlanModel, TaskRequest +from blueapi.service.model import PlanModel, TaskParametersV2, TaskRequestV2 FULL_PLAN = PlanModel( name="foobar", @@ -132,23 +132,38 @@ def test_plan_empty_fallback_help_text(client): @pytest.mark.parametrize( - "args,kwargs,params", + "plan_args", [ - p((1,), {}, {"one": 1}, id="required_as_positional"), - p((), {"one": 7}, {"one": 7}, id="required_as_keyword"), - p((1,), {"two": 23}, {"one": 1, "two": 23}, id="all_as_mixed_args_kwargs"), - p((1, 2), {}, {"one": 1, "two": 2}, id="all_as_positional"), - p((), {"one": 21, "two": 42}, {"one": 21, "two": 42}, id="all_as_keyword"), + pytest.param( + TaskParametersV2(args=(1,), kwargs={}), + id="required_as_positional", + ), + pytest.param( + TaskParametersV2(args=(), kwargs={"one": 7}), + id="required_as_keyword", + ), + pytest.param( + TaskParametersV2(args=(1,), kwargs={"two": 23}), + id="all_as_mixed_args_kwargs", + ), + p( + TaskParametersV2(args=(1, 2), kwargs={}), + id="all_as_positional", + ), + p( + TaskParametersV2(args=(), kwargs={"one": 21, "two": 42}), + id="all_as_keyword", + ), ], ) -def test_plan_param_mapping(args, kwargs, params): +def test_plan_param_mapping(plan_args): client = Mock() client.instrument_session = "cm12345-1" plan = Plan("foobar", FULL_PLAN, client) - plan(*args, **kwargs) + plan(*plan_args.args, **plan_args.kwargs) client.run_task.assert_called_once_with( - TaskRequest(name="foobar", instrument_session="cm12345-1", params=params) + TaskRequestV2(name="foobar", instrument_session="cm12345-1", params=plan_args) ) diff --git a/tests/unit_tests/client/test_rest.py b/tests/unit_tests/client/test_rest.py index 8c8c004cd..2fedd1ed9 100644 --- a/tests/unit_tests/client/test_rest.py +++ b/tests/unit_tests/client/test_rest.py @@ -35,16 +35,25 @@ EnvironmentResponse, PlanModel, TaskRequest, + TaskRequestV1, + TaskRequestV2, TaskResponse, TasksListResponse, WorkerTask, ) from blueapi.worker.event import WorkerState -from blueapi.worker.task import Task +from blueapi.worker.task import Task, TaskParametersV2 from blueapi.worker.task_worker import TrackableTask -TASK_REQUEST = TaskRequest( - name="foo", params={"one": "two"}, instrument_session="cm12345-1" +TASK_REQUEST_V1 = TaskRequestV1( + name="foo", + params={"one": "two"}, + instrument_session="cm12345-1", +) +TASK_REQUEST_V2 = TaskRequestV2( + name="foo", + params=TaskParametersV2(kwargs={"one": "two"}), + instrument_session="cm12345-1", ) @@ -89,17 +98,35 @@ def test_rest_error_code( rest.get_plans() -def test_create_task_serialization(): +@pytest.mark.parametrize( + "task_request", + [ + pytest.param( + TaskRequestV1( + name="demo", + instrument_session="cm12345-1", + params={"devices": [DeviceRef(name="foo", cache=Mock(), model=Mock())]}, + ), + id="v1", + ), + pytest.param( + TaskRequestV2( + name="demo", + instrument_session="cm12345-1", + params=TaskParametersV2( + kwargs={ + "devices": [DeviceRef(name="foo", cache=Mock(), model=Mock())] + } + ), + ), + id="v2", + ), + ], +) +def test_create_task_serialization(task_request: TaskRequest): rest = Mock(spec=BlueapiRestClient) - model = Mock() - model.name = "foo" - request = TaskRequest( - name="demo", - instrument_session="cm12345-1", - params={"devices": [DeviceRef(name="foo", cache=Mock(), model=model)]}, - ) - BlueapiRestClient.create_task(rest, request) + BlueapiRestClient.create_task(rest, task_request) rest._request_and_deserialize.assert_called_once_with( "/tasks", @@ -109,7 +136,7 @@ def test_create_task_serialization(): data={ "name": "demo", "instrument_session": "cm12345-1", - "params": {"devices": ["foo"]}, + "params": {"args": [], "kwargs": {"devices": ["foo"]}}, }, ) @@ -119,10 +146,10 @@ class CustomType: pass rest = Mock(spec=BlueapiRestClient) - request = TaskRequest( + request = TaskRequestV2( name="demo", instrument_session="cm12345-1", - params={"devices": [CustomType()]}, + params=TaskParametersV2(kwargs={"devices": [CustomType()]}), ) with pytest.raises(PydanticSerializationError, match="not serializable"): @@ -444,14 +471,17 @@ def test_get_missing_plan(rest: BlueapiRestClient): rest.get_plan("foo") +@pytest.mark.parametrize("task_request", [TASK_REQUEST_V1, TASK_REQUEST_V2]) @patch("blueapi.client.rest.connect") -def test_run_blocking(mock_connect: Mock, rest: BlueapiRestClient): +def test_run_blocking( + mock_connect: Mock, rest: BlueapiRestClient, task_request: TaskRequest +): ws = MagicMock() ws.__enter__.return_value.__iter__.return_value = iter( ['{"kind": "update", "data": {"name": "start", "doc":{}, "task_id":"t_uid"}}'] ) mock_connect.return_value = ws - conn = rest.run_blocking(TASK_REQUEST) + conn = rest.run_blocking(task_request) next(iter(conn)) mock_connect.assert_called_once_with( "ws://localhost:8000/api/v2/run_plan", @@ -460,6 +490,7 @@ def test_run_blocking(mock_connect: Mock, rest: BlueapiRestClient): ) +@pytest.mark.parametrize("task_request", [TASK_REQUEST_V1, TASK_REQUEST_V2]) @patch("blueapi.client.rest.connect") def test_run_blocking_auth( mock_connect: Mock, @@ -467,13 +498,14 @@ def test_run_blocking_auth( mock_authn_server: responses.RequestsMock, cached_valid_token: Path, # creates the cache file cached_valid_token_value: str, # the value from the file + task_request: TaskRequest, ): ws = MagicMock() ws.__enter__.return_value.__iter__.return_value = iter( ['{"kind": "update", "data": {"name": "start", "doc":{}, "task_id":"t_uid"}}'] ) mock_connect.return_value = ws - conn = rest_with_auth.run_blocking(TASK_REQUEST) + conn = rest_with_auth.run_blocking(task_request) next(iter(conn)) mock_connect.assert_called_once_with( "ws://localhost:8000/api/v2/run_plan", @@ -482,6 +514,7 @@ def test_run_blocking_auth( ) +@pytest.mark.parametrize("task_request", [TASK_REQUEST_V1, TASK_REQUEST_V2]) @pytest.mark.parametrize( "event,error,message", [ @@ -509,11 +542,12 @@ def test_run_blocking_errors( event: str, error: type[Exception], message: str, + task_request: TaskRequest, ): ws = MagicMock() ws.__enter__.return_value.__iter__.return_value = iter([event]) mock_connect.return_value = ws - conn = rest.run_blocking(TASK_REQUEST) + conn = rest.run_blocking(task_request) with pytest.raises(error, match=message): next(iter(conn)) mock_connect.assert_called_once_with( @@ -523,23 +557,30 @@ def test_run_blocking_errors( ) +@pytest.mark.parametrize("task_request", [TASK_REQUEST_V1, TASK_REQUEST_V2]) @pytest.mark.parametrize("status_code", [401, 403]) @patch("blueapi.client.rest.connect") def test_run_blocking_ws_failures( - mock_connect: Mock, rest: BlueapiRestClient, status_code: int + mock_connect: Mock, + rest: BlueapiRestClient, + status_code: int, + task_request: TaskRequest, ): mock_connect.side_effect = InvalidStatus( response=Response( status_code=status_code, reason_phrase="test_error", headers=Headers() ) ) - conn = rest.run_blocking(TASK_REQUEST) + conn = rest.run_blocking(task_request) with pytest.raises(UnauthorisedAccessError): next(iter(conn)) +@pytest.mark.parametrize("task_request", [TASK_REQUEST_V1, TASK_REQUEST_V2]) @patch("blueapi.client.rest.connect") -def test_run_blocking_unauthorised(mock_connect: Mock, rest: BlueapiRestClient): +def test_run_blocking_unauthorised( + mock_connect: Mock, rest: BlueapiRestClient, task_request: TaskRequest +): ws = MagicMock() event = """{ "kind":"unauthorized" @@ -547,11 +588,12 @@ def test_run_blocking_unauthorised(mock_connect: Mock, rest: BlueapiRestClient): ws.__enter__.return_value.__iter__.return_value = iter([event]) mock_connect.return_value = ws - conn = rest.run_blocking(TASK_REQUEST) + conn = rest.run_blocking(task_request) with pytest.raises(UnauthorisedAccessError): next(iter(conn)) +@pytest.mark.parametrize("task_request", [TASK_REQUEST_V1, TASK_REQUEST_V2]) @pytest.mark.parametrize( "conn_err,exp_err", [ @@ -575,8 +617,9 @@ def test_run_blocking_connect_error( rest: BlueapiRestClient, conn_err: Exception, exp_err: type[Exception], + task_request: TaskRequest, ): mock_connect.side_effect = conn_err - conn = rest.run_blocking(TASK_REQUEST) + conn = rest.run_blocking(task_request) with pytest.raises(exp_err): next(iter(conn)) diff --git a/tests/unit_tests/service/test_authorization.py b/tests/unit_tests/service/test_authorization.py index a2e602f21..0fddbde2f 100644 --- a/tests/unit_tests/service/test_authorization.py +++ b/tests/unit_tests/service/test_authorization.py @@ -13,7 +13,8 @@ submit_permission, validate_tiled_config, ) -from blueapi.service.model import TaskRequest +from blueapi.service.model import TaskRequestV2 +from blueapi.worker import TaskParametersV2 # Reusable client patch decorator patch_client_session = patch( @@ -194,7 +195,9 @@ async def test_user_client_can_submit_task(result, context: AbstractContextManag with context: await user_client.can_submit_task( - TaskRequest(name="foo", params={}, instrument_session="cm12345-1") + TaskRequestV2( + name="foo", params=TaskParametersV2(), instrument_session="cm12345-1" + ) ) opa.require_submit_task.assert_called_once_with("cm12345-1", "foo_bar") diff --git a/tests/unit_tests/service/test_interface.py b/tests/unit_tests/service/test_interface.py index 892c5ad2e..5bd163c7d 100644 --- a/tests/unit_tests/service/test_interface.py +++ b/tests/unit_tests/service/test_interface.py @@ -34,7 +34,8 @@ ProtocolInfo, PythonEnvironmentResponse, SourceInfo, - TaskRequest, + TaskRequestV1, + TaskRequestV2, WorkerTask, ) from blueapi.utils.invalid_config_error import InvalidConfigError @@ -46,7 +47,7 @@ WorkerEvent, WorkerState, ) -from blueapi.worker.task import Task +from blueapi.worker.task import Task, TaskParametersV2 from blueapi.worker.task_worker import TrackableTask FAKE_INSTRUMENT_SESSION = "cm12345-1" @@ -192,10 +193,26 @@ def test_get_device(context_mock: MagicMock): @patch("blueapi.service.interface.context") -def test_submit_task(context_mock: MagicMock): +def test_submit_task_v2(context_mock: MagicMock): context = BlueskyContext() context.register_plan(my_plan) - task = TaskRequest( + task = TaskRequestV2( + name="my_plan", + instrument_session=FAKE_INSTRUMENT_SESSION, + ) + context_mock.return_value = context + mock_uuid_value = "8dfbb9c2-7a15-47b6-bea8-b6b77c31d3d9" + with patch.object(uuid, "uuid4") as uuid_mock: + uuid_mock.return_value = uuid.UUID(mock_uuid_value) + task_uuid = interface.submit_task(task) + assert task_uuid == mock_uuid_value + + +@patch("blueapi.service.interface.context") +def test_submit_task_v1_converts_to_v2(context_mock: MagicMock): + context = BlueskyContext() + context.register_plan(my_plan) + task = TaskRequestV1( name="my_plan", instrument_session=FAKE_INSTRUMENT_SESSION, ) @@ -211,7 +228,7 @@ def test_submit_task(context_mock: MagicMock): def test_clear_task(context_mock: MagicMock): context = BlueskyContext() context.register_plan(my_plan) - task = TaskRequest( + task = TaskRequestV2( name="my_plan", instrument_session=FAKE_INSTRUMENT_SESSION, ) @@ -361,7 +378,7 @@ def test_get_task_by_id( context_mock.return_value = context task_id = interface.submit_task( - TaskRequest( + TaskRequestV2( name="my_plan", instrument_session=FAKE_INSTRUMENT_SESSION, ), @@ -385,7 +402,7 @@ def test_get_task_by_id( request_id=ANY, task=Task( name="my_plan", - params={}, + params=TaskParametersV2(), metadata=expected_metadata, ), is_complete=False, @@ -403,7 +420,7 @@ def test_submit_task_inserts_metadata(context_mock: MagicMock): metadata = {"foo": "bar"} task_id = interface.submit_task( - TaskRequest( + TaskRequestV2( name="my_plan", instrument_session=FAKE_INSTRUMENT_SESSION, ), @@ -417,7 +434,7 @@ def test_submit_task_inserts_metadata(context_mock: MagicMock): request_id=ANY, task=Task( name="my_plan", - params={}, + params=TaskParametersV2(), metadata=expected_metadata, ), is_complete=False, diff --git a/tests/unit_tests/service/test_protocol.py b/tests/unit_tests/service/test_protocol.py index 9178162e1..3acf6a47c 100644 --- a/tests/unit_tests/service/test_protocol.py +++ b/tests/unit_tests/service/test_protocol.py @@ -4,7 +4,7 @@ from pydantic import ValidationError from pydantic_core import InitErrorDetails -from blueapi.service.model import TaskRequest +from blueapi.service.model import TaskRequestV1, TaskRequestV2 from blueapi.service.protocol import ( Abort, ArgumentError, @@ -15,12 +15,13 @@ Resume, Submit, ) +from blueapi.worker import TaskParametersV2 @pytest.mark.parametrize( "src,res", [ - ( + pytest.param( """{ "kind": "submit", "task": { @@ -29,12 +30,35 @@ } }""", Submit( - task=TaskRequest(name="foo", params={}, instrument_session="cm12345-1") + task=TaskRequestV2( + name="foo", + params=TaskParametersV2(), + instrument_session="cm12345-1", + ) + ), + id="submit-v2", + ), + pytest.param( + """{ + "kind": "submit", + "task": { + "name": "foo", + "params": {"one": "two"}, + "instrument_session": "cm12345-1" + } + }""", + Submit( + task=TaskRequestV1( + name="foo", + params={"one": "two"}, + instrument_session="cm12345-1", + ) ), + id="submit-v1", ), - ('{"kind": "pause"}', Pause()), - ('{"kind": "resume"}', Resume()), - ('{"kind": "abort"}', Abort()), + pytest.param('{"kind": "pause"}', Pause(), id="pause"), + pytest.param('{"kind": "resume"}', Resume(), id="resume"), + pytest.param('{"kind": "abort"}', Abort(), id="abort"), ], ) def test_request_deserialization(src: str, res: Any): diff --git a/tests/unit_tests/service/test_rest_api.py b/tests/unit_tests/service/test_rest_api.py index a10d5ab3c..8bbbc9497 100644 --- a/tests/unit_tests/service/test_rest_api.py +++ b/tests/unit_tests/service/test_rest_api.py @@ -42,12 +42,12 @@ PythonEnvironmentResponse, SourceInfo, StateChangeRequest, - TaskRequest, + TaskRequestV2, WorkerTask, ) from blueapi.service.runner import WorkerDispatcher from blueapi.worker.event import TaskStatus, WorkerEvent, WorkerState -from blueapi.worker.task import Task +from blueapi.worker.task import Task, TaskParametersV2 from blueapi.worker.task_worker import TrackableTask @@ -58,7 +58,7 @@ class MockCountModel(BaseModel): ... FAKE_INSTRUMENT_SESSION = "cm12345-1" -SUBMIT_REQUEST = { +SUBMIT_REQUEST_V1 = { "kind": "submit", "task": { "name": "foo", @@ -67,6 +67,35 @@ class MockCountModel(BaseModel): ... }, } +SUBMIT_REQUEST_V2_ARGS = { + "kind": "submit", + "task": { + "name": "foo", + "params": {"args": ["two"], "kwargs": {}}, + "instrument_session": "cm12345-1", + }, +} + +SUBMIT_REQUEST_V2_KWARGS = { + "kind": "submit", + "task": { + "name": "foo", + "params": {"args": [], "kwargs": {"one": "two"}}, + "instrument_session": "cm12345-1", + }, +} + + +@pytest.fixture( + params=[ + pytest.param(SUBMIT_REQUEST_V1, id="v1"), + pytest.param(SUBMIT_REQUEST_V2_ARGS, id="v2_args"), + pytest.param(SUBMIT_REQUEST_V2_KWARGS, id="v2_kwargs"), + ] +) +def submit_request(request: pytest.FixtureRequest): + return request.param + @pytest.fixture def mock_runner() -> Mock: @@ -170,9 +199,9 @@ def test_rest_config_with_cors( client_with_cors: TestClient, mock_runner: Mock, ): - task = TaskRequest( + task = TaskRequestV2( name="my-plan", - params={"id": "x"}, + params=TaskParametersV2(kwargs={"id": "x"}), instrument_session=FAKE_INSTRUMENT_SESSION, ) task_id = "f8424be3-203c-494e-b22f-219933b4fa67" @@ -285,9 +314,9 @@ def test_get_non_existent_device_by_name(mock_runner: Mock, client: TestClient) def test_create_task(mock_runner: Mock, client: TestClient) -> None: - task = TaskRequest( + task = TaskRequestV2( name="count", - params={"detectors": ["x"]}, + params=TaskParametersV2(kwargs={"detectors": ["x"]}), instrument_session=FAKE_INSTRUMENT_SESSION, ) task_id = str(uuid.uuid4()) @@ -306,7 +335,11 @@ def test_submit_task_requires_permission( mock_opa_client: Mock, access_token: str, ): - task = TaskRequest(name="sleep", params={"time": 2}, instrument_session="cm12345-2") + task = TaskRequestV2( + name="sleep", + params=TaskParametersV2(kwargs={"time": 2}), + instrument_session="cm12345-2", + ) client_with_opa.headers["Authorization"] = f"Bearer {access_token}" mock_opa_client.can_submit_task.side_effect = HTTPException(status_code=403) mock_runner.run.side_effect = RuntimeError("Task should not be submitted") @@ -321,9 +354,9 @@ def test_create_task_inserts_auth_metadata( mock_runner: Mock, client_with_auth: TestClient, ) -> None: - task = TaskRequest( + task = TaskRequestV2( name="count", - params={"detectors": ["x"]}, + params=TaskParametersV2(kwargs={"detectors": ["x"]}), instrument_session=FAKE_INSTRUMENT_SESSION, ) client_with_auth.follow_redirects = False @@ -395,7 +428,10 @@ def test_put_plan_fails_if_not_idle(mock_runner: Mock, client: TestClient) -> No def test_get_tasks(mock_runner: Mock, client: TestClient) -> None: tasks = [ - TrackableTask(task_id="0", task=Task(name="sleep", params={"time": 0.0})), + TrackableTask( + task_id="0", + task=Task(name="sleep", params=TaskParametersV2(kwargs={"time": 0.0})), + ), TrackableTask( task_id="1", task=Task(name="first_task"), @@ -418,7 +454,7 @@ def test_get_tasks(mock_runner: Mock, client: TestClient) -> None: "request_id": None, "task": { "name": "sleep", - "params": {"time": 0.0}, + "params": {"args": [], "kwargs": {"time": 0.0}}, "metadata": {}, }, "outcome": None, @@ -431,7 +467,7 @@ def test_get_tasks(mock_runner: Mock, client: TestClient) -> None: "request_id": None, "task": { "name": "first_task", - "params": {}, + "params": {"args": [], "kwargs": {}}, "metadata": {}, }, "outcome": None, @@ -463,7 +499,7 @@ def test_get_tasks_by_status(mock_runner: Mock, client: TestClient) -> None: "request_id": None, "task": { "name": "third_task", - "params": {}, + "params": {"args": [], "kwargs": {}}, "metadata": {}, }, "outcome": None, @@ -627,7 +663,7 @@ def test_get_task(mock_runner: Mock, client: TestClient): "request_id": None, "task": { "name": "third_task", - "params": {}, + "params": {"args": [], "kwargs": {}}, "metadata": { "foo": "bar", }, @@ -674,7 +710,7 @@ def test_get_all_tasks(mock_runner: Mock, client: TestClient): "task_id": task_id, "task": { "name": "third_task", - "params": {}, + "params": {"args": [], "kwargs": {}}, "metadata": {}, }, "is_complete": False, @@ -1010,7 +1046,9 @@ def test_logout_when_oidc_config_invalid( assert response.status_code == status.HTTP_205_RESET_CONTENT -async def test_websocket_run_plan(mock_runner: Mock, client: TestClient): +async def test_websocket_run_plan( + mock_runner: Mock, client: TestClient, submit_request +): mock_runner.run.side_effect = lambda req, *a, **kw: { interface.get_active_task: None, interface.submit_task: "task_id", @@ -1040,7 +1078,7 @@ async def test_websocket_run_plan(mock_runner: Mock, client: TestClient): ) with client.websocket_connect("/api/v2/run_plan") as ws: - ws.send_json(SUBMIT_REQUEST) + ws.send_json(submit_request) assert ws.receive_json() == { "kind": "update", "data": { @@ -1113,10 +1151,11 @@ def test_websocket_run_plan_submit_error( reason, mock_runner: Mock, client: TestClient, + submit_request, ): mock_runner.run.side_effect = exc with client.websocket_connect("/api/v2/run_plan") as ws: - ws.send_json(SUBMIT_REQUEST) + ws.send_json(submit_request) assert ws.receive_json() == err_message with pytest.raises(WebSocketDisconnect) as disco: ws.receive_text() @@ -1124,13 +1163,15 @@ def test_websocket_run_plan_submit_error( assert disco.value.reason == reason -def test_websocket_run_plan_server_busy(mock_runner: Mock, client: TestClient): +def test_websocket_run_plan_server_busy( + mock_runner: Mock, client: TestClient, submit_request +): mock_runner.run.side_effect = [ "task_id", # submit_task Mock(name="active_task", is_complete=False), ] with client.websocket_connect("/api/v2/run_plan") as ws: - ws.send_json(SUBMIT_REQUEST) + ws.send_json(submit_request) assert ws.receive_json() == {"kind": "busy"} with pytest.raises(WebSocketDisconnect) as disco: ws.receive_text() @@ -1139,7 +1180,9 @@ def test_websocket_run_plan_server_busy(mock_runner: Mock, client: TestClient): pass -def test_websocket_run_plan_unrelated_events(mock_runner: Mock, client: TestClient): +def test_websocket_run_plan_unrelated_events( + mock_runner: Mock, client: TestClient, submit_request +): mock_runner.run.side_effect = lambda req, *a, **kw: { interface.get_active_task: None, interface.submit_task: "task_id", @@ -1169,7 +1212,7 @@ def test_websocket_run_plan_unrelated_events(mock_runner: Mock, client: TestClie ) with client.websocket_connect("/api/v2/run_plan") as ws: - ws.send_json(SUBMIT_REQUEST) + ws.send_json(submit_request) # first event is not sent assert ws.receive_json() == { "kind": "update", @@ -1195,7 +1238,7 @@ def test_websocket_run_plan_unrelated_events(mock_runner: Mock, client: TestClie def test_websocket_run_plan_client_disconnect_cancels( - mock_runner: Mock, client: TestClient + mock_runner: Mock, client: TestClient, submit_request ): mock_runner.run.side_effect = ["task_id", None, None, None] mock_runner.event_pipe.return_value = contextlib.nullcontext( @@ -1227,7 +1270,7 @@ async def local_send(message: Message): cast(FastAPI, client.app).add_middleware(Disconnector) with client.websocket_connect("/api/v2/run_plan") as ws: - ws.send_json(SUBMIT_REQUEST) + ws.send_json(submit_request) mock_runner.run.assert_called_with( interface.cancel_active_task, failure=True, reason="Client disconnected" ) @@ -1249,6 +1292,7 @@ def test_websocket_run_plan_needs_auth_token( def test_websocket_run_plan_needs_permission( mock_runner: Mock, client_with_opa: TestClient, + submit_request, mock_opa_client: Mock, access_token: str, ): @@ -1256,7 +1300,7 @@ def test_websocket_run_plan_needs_permission( mock_opa_client.can_submit_task.side_effect = HTTPException(status_code=403) mock_runner.run.side_effect = RuntimeError("Task should not be submitted") with client_with_opa.websocket_connect("/api/v2/run_plan") as ws: - ws.send_json(SUBMIT_REQUEST) + ws.send_json(submit_request) assert ws.receive_json() == {"kind": "unauthorized"} with pytest.raises(WebSocketDisconnect) as discon: ws.receive_text() diff --git a/tests/unit_tests/worker/test_task_worker.py b/tests/unit_tests/worker/test_task_worker.py index f6fe50c0d..3c982a5bf 100644 --- a/tests/unit_tests/worker/test_task_worker.py +++ b/tests/unit_tests/worker/test_task_worker.py @@ -25,6 +25,7 @@ from blueapi.utils.base_model import BlueapiBaseModel from blueapi.worker import ( Task, + TaskParametersV2, TaskStatus, TaskWorker, TrackableTask, @@ -35,16 +36,16 @@ ) from blueapi.worker.event import TaskResult, TaskStatusEnum -_SIMPLE_TASK = Task(name="sleep", params={"time": 0.0}) -_LONG_TASK = Task(name="sleep", params={"time": 1.0}) +_SIMPLE_TASK = Task(name="sleep", params=TaskParametersV2(kwargs={"time": 0.0})) +_LONG_TASK = Task(name="sleep", params=TaskParametersV2(kwargs={"time": 1.0})) _INDEFINITE_TASK = Task( name="set_absolute", - params={"movable": "fake_device", "value": 4.0}, + params=TaskParametersV2(kwargs={"movable": "fake_device", "value": 4.0}), ) -_FAILING_TASK = Task(name="failing_plan", params={}) +_FAILING_TASK = Task(name="failing_plan", params=TaskParametersV2()) _TASK_WITH_METADATA = Task( name="sleep", - params={"time": 0.0}, + params=TaskParametersV2(kwargs={"time": 0.0}), metadata={ "foo": "bar", "baz": 0, @@ -426,6 +427,17 @@ def test_metadata_passed_to_context(context: BlueskyContext): assert metadata in context.run_engine.md.items() +def test_do_task_with_invalid_parameters_raises_validation_error( + context: BlueskyContext, +): + context.run_engine = Mock() + task = Task(name="sleep", params=TaskParametersV2()) + # Required as this is the error that blueapi/service/main:submit_task is expecting + # to catch. + with pytest.raises(pydantic.ValidationError): + task.do_task(context) + + # # Worker helpers # @@ -509,7 +521,9 @@ def assert_running_count_plan_produces_ordered_worker_and_data_events( task: Task | None = None, timeout: float = 5.0, ) -> None: - default_task = Task(name="count", params={"detectors": ["motor"], "num": 1}) + default_task = Task( + name="count", params=TaskParametersV2(kwargs={"detectors": ["motor"], "num": 1}) + ) task = task or default_task event_streams: list[EventStream[Any, int]] = [ @@ -620,7 +634,10 @@ def test_get_tasks(worker: TaskWorker, status, expected_task_ids): "task1": TrackableTask( task_id="task1", task=Task( - name="set_absolute", params={"movable": "fake_device", "value": 4.0} + name="set_absolute", + params=TaskParametersV2( + kwargs={"movable": "fake_device", "value": 4.0} + ), ), is_complete=False, is_pending=False, @@ -628,7 +645,10 @@ def test_get_tasks(worker: TaskWorker, status, expected_task_ids): "task2": TrackableTask( task_id="task2", task=Task( - name="set_absolute", params={"movable": "fake_device", "value": 4.0} + name="set_absolute", + params=TaskParametersV2( + kwargs={"movable": "fake_device", "value": 4.0} + ), ), is_complete=False, is_pending=True, @@ -638,7 +658,10 @@ def test_get_tasks(worker: TaskWorker, status, expected_task_ids): "task3": TrackableTask( task_id="task3", task=Task( - name="set_absolute", params={"movable": "fake_device", "value": 4.0} + name="set_absolute", + params=TaskParametersV2( + kwargs={"movable": "fake_device", "value": 4.0} + ), ), is_complete=True, is_pending=False, @@ -724,8 +747,8 @@ def injected_device_plan( yield from () context.register_plan(injected_device_plan) - params = Task(name="injected_device_plan").prepare_params(context) - assert params["dev"] == fake_device + args, kwargs = Task(name="injected_device_plan").prepare_params(context) + assert kwargs["dev"] == fake_device def test_injected_devices_plan_model( @@ -840,9 +863,9 @@ def test_injected_composite_devices_are_found( context: BlueskyContext, ): context.register_plan(injected_device_plan) - params = Task(name="injected_device_plan").prepare_params(context) - assert params["composite"].fake_device == fake_device - assert params["composite"].second_fake_device == second_fake_device + args, kwargs = Task(name="injected_device_plan").prepare_params(context) + assert kwargs["composite"].fake_device == fake_device + assert kwargs["composite"].second_fake_device == second_fake_device def test_injected_composite_devices_plan_model( @@ -861,9 +884,9 @@ def test_injected_composite_with_pydantic_dataclass( second_fake_device: FakeDevice, ): context.register_plan(injected_dataclass_device_plan) - params = Task(name="injected_dataclass_device_plan").prepare_params(context) - assert params["composite"].fake_device == fake_device - assert params["composite"].second_fake_device == second_fake_device + args, kwargs = Task(name="injected_dataclass_device_plan").prepare_params(context) + assert kwargs["composite"].fake_device == fake_device + assert kwargs["composite"].second_fake_device == second_fake_device def test_injected_composite_with_standard_dataclass( @@ -872,11 +895,11 @@ def test_injected_composite_with_standard_dataclass( second_fake_device: FakeDevice, ): context.register_plan(injected_standard_dataclass_device_plan) - params = Task(name="injected_standard_dataclass_device_plan").prepare_params( + args, kwargs = Task(name="injected_standard_dataclass_device_plan").prepare_params( context ) - assert params["composite"].fake_device == fake_device - assert params["composite"].second_fake_device == second_fake_device + assert kwargs["composite"].fake_device == fake_device + assert kwargs["composite"].second_fake_device == second_fake_device def test_plan_module_with_composite_devices_can_be_loaded_before_device_module( @@ -887,9 +910,11 @@ def test_plan_module_with_composite_devices_can_be_loaded_before_device_module( context_without_devices.register_plan(injected_device_plan) context_without_devices.register_device(fake_device) context_without_devices.register_device(second_fake_device) - params = Task(name="injected_device_plan").prepare_params(context_without_devices) - assert params["composite"].fake_device == fake_device - assert params["composite"].second_fake_device == second_fake_device + args, kwargs = Task(name="injected_device_plan").prepare_params( + context_without_devices + ) + assert kwargs["composite"].fake_device == fake_device + assert kwargs["composite"].second_fake_device == second_fake_device @pytest.mark.parametrize( @@ -949,3 +974,175 @@ def test_worker_event_task_id(): def test_worker_event_no_task_id(): event = WorkerEvent(state=WorkerState.IDLE, task_status=None) assert event.task_id is None + + +def test_task_worker_passes_positional_args( + context: BlueskyContext, +) -> None: + + def positional_plan(value: int) -> MsgGenerator: + yield from () + + context.register_plan(positional_plan) + + task = Task(name="positional_plan", params=TaskParametersV2(args=[42])) + args, kwargs = task.prepare_params(context) + + assert args == [42] + assert kwargs == {} + + +def test_task_worker_passes_multiple_positional_args( + context: BlueskyContext, +) -> None: + + def positional_plan(first: int, second: int) -> MsgGenerator: + yield from () + + context.register_plan(positional_plan) + task = Task(name="positional_plan", params=TaskParametersV2(args=[1, 2])) + args, kwargs = task.prepare_params(context) + + assert args == [1, 2] + assert kwargs == {} + + +def test_task_worker_passes_positional_and_keyword_args( + context: BlueskyContext, +) -> None: + + def mixed_plan( + first: int, + second: int, + *, + third: int, + ) -> MsgGenerator: + yield from () + + context.register_plan(mixed_plan) + task = Task( + name="mixed_plan", + params=TaskParametersV2(args=[1, 2], kwargs={"third": 3}), + ) + args, kwargs = task.prepare_params(context) + assert args == [1, 2] + assert kwargs == {"third": 3} + + +def test_task_worker_resolves_positional_device( + context: BlueskyContext, + fake_device: FakeDevice, +) -> None: + def positional_device_plan(device: FakeDevice) -> MsgGenerator: + yield from () + + context.register_plan(positional_device_plan) + + task = Task( + name="positional_device_plan", params=TaskParametersV2(args=[fake_device.name]) + ) + args, kwargs = task.prepare_params(context) + assert args == [fake_device] + assert kwargs == {} + + +def test_positional_only_plan(fake_device: FakeDevice) -> None: + def positional_only_plan(val: int, /, device: FakeDevice) -> MsgGenerator: + yield from () + + context = BlueskyContext() + context.register_device(fake_device) + context.register_plan(positional_only_plan) + + task = Task( + name="positional_only_plan", + params=TaskParametersV2( + args=[1], + kwargs={"device": fake_device.name}, + ), + ) + + args, kwargs = task.prepare_params(context) + assert args == [1] + assert kwargs == {"device": fake_device} + + +def test_var_positional_plan(fake_device: FakeDevice) -> None: + def positional_device_plan(val: int, *device: FakeDevice) -> MsgGenerator: + yield from () + + context = BlueskyContext() + context.register_device(fake_device) + context.register_plan(positional_device_plan) + + task = Task( + name="positional_device_plan", + params=TaskParametersV2(args=[1, fake_device.name, fake_device.name]), + ) + args, kwargs = task.prepare_params(context) + assert args == [1, fake_device, fake_device] + assert kwargs == {} + + +def test_positional_only_and_var_positional_plan( + fake_device: FakeDevice, +) -> None: + def positional_device_plan(val: int, /, *device: FakeDevice) -> MsgGenerator: + yield from () + + context = BlueskyContext() + context.register_device(fake_device) + context.register_plan(positional_device_plan) + + task = Task( + name="positional_device_plan", + params=TaskParametersV2( + args=[1, fake_device.name, fake_device.name], + ), + ) + + args, kwargs = task.prepare_params(context) + assert args == [1, fake_device, fake_device] + assert kwargs == {} + + +def test_var_keyword_plan(fake_device: FakeDevice) -> None: + def keyword_device_plan(val: int, **devices: FakeDevice) -> MsgGenerator: + yield from () + + context = BlueskyContext() + context.register_device(fake_device) + context.register_plan(keyword_device_plan) + + task = Task( + name="keyword_device_plan", + params=TaskParametersV2( + args=[1], + kwargs={"first": fake_device.name, "second": fake_device.name}, + ), + ) + args, kwargs = task.prepare_params(context) + assert args == [1] + assert kwargs == {"first": fake_device, "second": fake_device} + + +def test_positional_only_and_var_keyword_plan( + fake_device: FakeDevice, +) -> None: + def positional_only_plan(val: int, /, **devices: FakeDevice) -> MsgGenerator: + yield from () + + context = BlueskyContext() + context.register_device(fake_device) + context.register_plan(positional_only_plan) + + task = Task( + name="positional_only_plan", + params=TaskParametersV2( + args=[1], + kwargs={"first": fake_device.name, "second": fake_device.name}, + ), + ) + args, kwargs = task.prepare_params(context) + assert args == [1] + assert kwargs == {"first": fake_device, "second": fake_device}