diff --git a/docs/examples/add_tasks_using_client.py b/docs/examples/add_tasks_using_client.py new file mode 100644 index 0000000..e62c700 --- /dev/null +++ b/docs/examples/add_tasks_using_client.py @@ -0,0 +1,46 @@ +from blueapi.service.model import TaskRequest + +from daq_queuing_service.client.client import QueueClient +from daq_queuing_service.task_queue.queue import QueueState + +# Example of how to add tasks to the queue using the Python client +# Need a local queue running on port 8001 for this to work + + +def main(): + client = QueueClient("http://127.0.0.1:8001") + + # Pause the queue + client.update_queue_state(QueueState(paused=True)) + + # Clear everything in the queue + client.cancel_all_tasks() + + # Add 10 sleep plans to the queue + client.add_tasks_to_queue( + [ + TaskRequest( + name="sleep", instrument_session="cm44163-3", params={"time": 1} + ) + for _ in range(10) + ], + ) + + # Add a longer sleep to the start of the queue + client.add_tasks_to_queue( + [TaskRequest(name="sleep", instrument_session="cm44163-3", params={"time": 5})], + position=0, + ) + + # Print the current tasks in the queue + tasks = client.get_queued_tasks() + for task in tasks: + assert isinstance(task.experiment, TaskRequest) + print( + f"Task {task.position}: {task.experiment.name} for " + + f"{task.experiment.params['time']} seconds" + ) + + +if __name__ == "__main__": + main() diff --git a/src/daq_queuing_service/api/api.py b/src/daq_queuing_service/api/api.py index da3aec3..efdeabf 100644 --- a/src/daq_queuing_service/api/api.py +++ b/src/daq_queuing_service/api/api.py @@ -45,11 +45,11 @@ def create_api_router( router = APIRouter() @router.get("/healthz") - async def healthz(): + async def healthz() -> Response: return Response() @router.get("/") - def read_root(request: Request): + def read_root(request: Request) -> str: base_url = str(request.base_url) return ( f"Welcome to the daq queuing service. Visit {base_url}docs for Uvicorn API." diff --git a/src/daq_queuing_service/api/errors.py b/src/daq_queuing_service/api/errors.py index ddebaf0..93b249b 100644 --- a/src/daq_queuing_service/api/errors.py +++ b/src/daq_queuing_service/api/errors.py @@ -1,5 +1,6 @@ from fastapi import FastAPI, Request from fastapi.responses import JSONResponse +from pydantic import BaseModel from daq_queuing_service.task_queue.queue_utils import ( NegativePositionError, @@ -9,7 +10,11 @@ TaskNotInQueueError, ) + # pyright: reportUnusedFunction=false +class ErrorContent(BaseModel): + error: str + message: str def register_exception_handlers(app: FastAPI): @@ -19,14 +24,18 @@ async def task_in_progress_handler( ): return JSONResponse( status_code=409, - content={"error": "task_in_progress", "message": str(exception)}, + content=ErrorContent( + error="task_in_progress", message=str(exception) + ).model_dump(), ) @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)}, + content=ErrorContent( + error="task_not_found", message=str(exception) + ).model_dump(), ) @app.exception_handler(TaskNotInQueueError) @@ -35,7 +44,9 @@ async def task_not_in_queue_handler( ): return JSONResponse( status_code=409, - content={"error": "task_not_in_queue", "message": str(exception)}, + content=ErrorContent( + error="task_not_in_queue", message=str(exception) + ).model_dump(), ) @app.exception_handler(NegativePositionError) @@ -44,12 +55,16 @@ async def negative_position_handler( ): return JSONResponse( status_code=400, - content={"error": "negative_position", "message": str(exception)}, + content=ErrorContent( + error="negative_position", message=str(exception) + ).model_dump(), ) @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)}, + content=ErrorContent( + error="queue_error", message=str(exception) + ).model_dump(), ) diff --git a/src/daq_queuing_service/blueapi_interaction/clients.py b/src/daq_queuing_service/blueapi_interaction/clients.py index 3591191..7677441 100644 --- a/src/daq_queuing_service/blueapi_interaction/clients.py +++ b/src/daq_queuing_service/blueapi_interaction/clients.py @@ -1,5 +1,3 @@ -from unittest.mock import MagicMock - from blueapi.client import BlueapiClient from blueapi.client.event_bus import EventBusClient from blueapi.client.rest import BlueapiRestClient @@ -10,9 +8,6 @@ def get_blueapi_clients(blueapi_config: ApplicationConfig) -> BlueapiClient: - if not blueapi_config.oidc: - blueapi_config.oidc = MagicMock() - blueapi_rest_client = BlueapiRestClient( config=blueapi_config.api, # Waiting on https://github.com/DiamondLightSource/blueapi/pull/1553 diff --git a/src/daq_queuing_service/client/client.py b/src/daq_queuing_service/client/client.py new file mode 100644 index 0000000..a69c1c1 --- /dev/null +++ b/src/daq_queuing_service/client/client.py @@ -0,0 +1,147 @@ +from collections.abc import Mapping +from typing import Any, TypeVar + +import requests +from blueapi.service.model import TaskRequest +from pydantic import HttpUrl, TypeAdapter, ValidationError + +from daq_queuing_service.api.api import TaskCancelRequest +from daq_queuing_service.api.errors import ErrorContent +from daq_queuing_service.app._config import AppConfig +from daq_queuing_service.blueapi_interaction.blueapi_call import BlueapiCallResponse +from daq_queuing_service.task import Experiment, TaskWithPosition +from daq_queuing_service.task_queue.queue import QueueState + +T = TypeVar("T") + + +class QueueClient: + def __init__(self, url: str): + self._url = HttpUrl(url) + self._pool = requests.Session() + + def _request( + self, + suffix: str, + method: str = "GET", + data: Any = None, + params: Mapping[str, Any] | None = None, + ): + url = self._url.unicode_string().removesuffix("/") + suffix + response = self._pool.request( + method, + url, + json=data, + params=params, + ) + return response + + def _request_expect_error( + self, + suffix: str, + target_type: type[T], + method: str = "GET", + data: Any = None, + params: Mapping[str, Any] | None = None, + ): + response = self._request(suffix, method, data, params) + try: + return TypeAdapter(target_type).validate_python(response.json()) + except ValidationError: + return ErrorContent.model_validate(response.json()) + + def _request_expect_none( + self, + suffix: str, + target_type: type[T], + method: str = "GET", + data: Any = None, + params: Mapping[str, Any] | None = None, + ) -> T | None: + response = self._request(suffix, method, data, params) + if response.json() is None: + return + return TypeAdapter(target_type).validate_python(response.json()) + + def _request_and_validate( + self, + suffix: str, + target_type: type[T], + method: str = "GET", + data: Any = None, + params: Mapping[str, Any] | None = None, + ) -> T: + response = self._request(suffix, method, data, params) + return TypeAdapter(target_type).validate_python(response.json()) + + def healthz(self): + return self._request_and_validate("/healthz", str) + + def get_config(self): + return self._request_and_validate("/config", AppConfig) + + def get_queue_state(self) -> QueueState: + return self._request_and_validate("/queue/state", QueueState) + + def update_queue_state(self, new_state: QueueState) -> QueueState: + return self._request_and_validate( + "/queue/state", QueueState, method="PATCH", data=new_state.model_dump() + ) + + def get_queued_tasks(self) -> list[TaskWithPosition]: + return self._request_and_validate("/queue", list[TaskWithPosition]) + + def add_tasks_to_queue( + self, + experiments: list[Experiment | TaskRequest], + position: int | None = None, + ) -> list[str] | ErrorContent: + return self._request_expect_error( + "/queue", + list[str], + method="POST", + data=[experiment.model_dump() for experiment in experiments], + params={"position": position}, + ) + + def move_task(self, task_id: str, new_position: int) -> int | ErrorContent: + return self._request_expect_error( + "/queue/move", + int, + method="POST", + params={"task_id": task_id, "new_position": new_position}, + ) + + def cancel_tasks( + self, task_ids: list[str] + ) -> list[TaskWithPosition] | ErrorContent: + return self._request_expect_error( + "/queue/tasks", + list[TaskWithPosition], + method="DELETE", + data=TaskCancelRequest(task_ids=task_ids).model_dump(), + ) + + def cancel_all_tasks(self) -> list[TaskWithPosition]: + return self._request_and_validate("/queue", list[TaskWithPosition], "DELETE") + + def get_task_by_position(self, position: int) -> TaskWithPosition | None: + return self._request_expect_none(f"/queue/{position}", TaskWithPosition) + + def get_all_tasks(self) -> list[TaskWithPosition]: + return self._request_and_validate("/tasks", list[TaskWithPosition]) + + def get_task_by_id(self, task_id: str) -> TaskWithPosition | ErrorContent: + return self._request_expect_error(f"/tasks/{task_id}", TaskWithPosition) + + def get_completed_tasks(self) -> list[TaskWithPosition]: + return self._request_and_validate("/history", list[TaskWithPosition]) + + def clear_history(self): + return self._request_and_validate("/history", str, "DELETE") + + def get_call_queue(self) -> list[BlueapiCallResponse]: + return self._request_and_validate("/call_queue", list[BlueapiCallResponse]) + + def get_call_history(self) -> list[BlueapiCallResponse]: + return self._request_and_validate("/call_history", list[BlueapiCallResponse]) diff --git a/src/daq_queuing_service/worker/worker.py b/src/daq_queuing_service/worker/worker.py index 11ed2bf..bae0260 100644 --- a/src/daq_queuing_service/worker/worker.py +++ b/src/daq_queuing_service/worker/worker.py @@ -80,7 +80,7 @@ async def _process_call(self, call: BlueapiCall): match task_status.result: case TaskResult(): LOGGER.debug( - f"Call {call} completed succesfully: {task_status.result}" + f"Call {call} completed successfully: {task_status.result}" ) await self._queue.complete_call(call, task_status.result) case TaskError(): @@ -115,8 +115,7 @@ async def _handle_run_task_error( error: InvalidParametersError | UnknownPlanError | ServiceUnavailableError - | BlueskyRemoteControlError - | ServiceUnavailableError, + | BlueskyRemoteControlError, ): match error: case InvalidParametersError(): diff --git a/tests/unit_tests/client/test_client.py b/tests/unit_tests/client/test_client.py new file mode 100644 index 0000000..72e0785 --- /dev/null +++ b/tests/unit_tests/client/test_client.py @@ -0,0 +1,63 @@ +from typing import Any +from unittest.mock import MagicMock, _Call, call, patch + +import pytest + +from daq_queuing_service.api.api import create_api_router +from daq_queuing_service.client.client import QueueClient + + +def test_all_endpoints_have_a_client_function(): + exceptions = ["read_root", "stream_events"] + + router = create_api_router(MagicMock(), MagicMock(), MagicMock()) + endpoints: list[str] = [route.endpoint.__name__ for route in router.routes] # type: ignore + + not_in_client = [ + endpoint + for endpoint in endpoints + if endpoint not in list(QueueClient.__dict__.keys()) + exceptions + ] + + assert not not_in_client, ( + f"Found endpoints not covered by a client function: {not_in_client}" + ) + + +@pytest.mark.parametrize( + "url, kwargs, expected_request_call", + [ + ( + "https://google.com/", + {"suffix": "/queue"}, + call("GET", "https://google.com/queue", json=None, params=None), + ), + ( + "http://google.com", + { + "suffix": "/queue", + "method": "POST", + "data": {"test": "data"}, + "params": {"test": "params"}, + }, + call( + "POST", + "http://google.com/queue", + json={"test": "data"}, + params={"test": "params"}, + ), + ), + ], +) +def test__request_makes_request_with_expected_arguments( + url: str, kwargs: dict[str, Any], expected_request_call: _Call +): + client = QueueClient(url=url) + + with patch( + "daq_queuing_service.client.client.requests.Session.request" + ) as mock_request: + client._request(**kwargs) + + mock_request.assert_called_once() + mock_request.assert_has_calls([expected_request_call])