fix: get tasks endpoint to a enum (#1578)

Currently get_tasks takes a string to filter on task_status, but it
should be a enum.

Author: ZohebShaikh

docs/reference/openapi.yaml

@@ -333,6 +333,14 @@ components:
       - type
       title: TaskResult
       type: object
+    TaskStatusEnum:
+      enum:
+      - PENDING
+      - COMPLETE
+      - ERROR
+      - RUNNING
+      title: TaskStatusEnum
+      type: string
     TasksListResponse:
       additionalProperties: false
       description: Diagnostic information on the tasks
@@ -441,7 +449,7 @@ info:
     name: Apache 2.0
     url: https://www.apache.org/licenses/LICENSE-2.0.html
   title: BlueAPI Control
-  version: 1.4.0
+  version: 1.4.1
 openapi: 3.1.0
 paths:
   /api/v1/devices:
@@ -599,15 +607,17 @@ paths:
     get:
       description: 'Retrieve tasks based on their status.
 
-        The status of a newly created task is ''unstarted''.'
+        The status of a newly created task is PENDING.'
       operationId: get_tasks_api_v1_tasks_get
       parameters:
       - in: query
         name: task_status
         required: false
         schema:
+          anyOf:
+          - $ref: '#/components/schemas/TaskStatusEnum'
+          - type: 'null'
           title: Task Status
-          type: string
       responses:
         '200':
           content:
@@ -997,15 +1007,17 @@ paths:
       deprecated: true
       description: 'Retrieve tasks based on their status.
 
-        The status of a newly created task is ''unstarted''.'
+        The status of a newly created task is PENDING.'
       operationId: get_tasks_tasks_get
       parameters:
       - in: query
         name: task_status
         required: false
         schema:
+          anyOf:
+          - $ref: '#/components/schemas/TaskStatusEnum'
+          - type: 'null'
           title: Task Status
-          type: string
       responses:
         '200':
           content:

src/blueapi/config.py

@@ -322,7 +322,7 @@ class ApplicationConfig(BlueapiBaseModel):
     """
 
     #: API version to publish in OpenAPI schema
-    REST_API_VERSION: ClassVar[str] = "1.4.0"
+    REST_API_VERSION: ClassVar[str] = "1.4.1"
 
     LICENSE_INFO: ClassVar[dict[str, str]] = {
         "name": "Apache 2.0",

src/blueapi/service/interface.py

@@ -233,11 +233,6 @@ def remove_callback_when_task_finished(
     return task
 
 
-def get_tasks_by_status(status: TaskStatusEnum) -> list[TrackableTask]:
-    """Retrieve a list of tasks based on their status."""
-    return worker().get_tasks_by_status(status)
-
-
 def get_active_task() -> TrackableTask | None:
     """Task the worker is currently running"""
     return worker().get_active_task()
@@ -264,10 +259,10 @@ def cancel_active_task(failure: bool, reason: str | None) -> str:
     return worker().cancel_active_task(failure, reason)
 
 
-def get_tasks() -> list[TrackableTask]:
-    """Return a list of all tasks on the worker,
-    any one of which can be triggered with begin_task"""
-    return worker().get_tasks()
+def get_tasks(status: TaskStatusEnum | None = None) -> list[TrackableTask]:
+    """Retrieve a list of tasks based on their status.
+    Return a list of all tasks on the worker if status is None"""
+    return worker().get_tasks(status)
 
 
 def get_task_by_id(task_id: str) -> TrackableTask | None:

src/blueapi/service/main.py

@@ -27,7 +27,6 @@
 from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 from opentelemetry.trace import get_tracer_provider
 from pydantic import ValidationError
-from pydantic.json_schema import SkipJsonSchema
 from starlette.responses import JSONResponse
 from super_state_machine.errors import TransitionError
 
@@ -317,38 +316,18 @@ def delete_submitted_task(
     return TaskResponse(task_id=runner.run(interface.clear_task, task_id))
 
 
-@start_as_current_span(TRACER, "v")
-def validate_task_status(v: str) -> TaskStatusEnum:
-    v_upper = v.upper()
-    if v_upper not in TaskStatusEnum.__members__:
-        raise ValueError("Invalid status query parameter")
-    return TaskStatusEnum(v_upper)
-
-
 @secure_router_v1.get("/tasks", status_code=status.HTTP_200_OK, tags=[Tag.TASK])
 @secure_router.get("/tasks", status_code=status.HTTP_200_OK, tags=[Tag.TASK])
 @start_as_current_span(TRACER)
 def get_tasks(
     runner: Annotated[WorkerDispatcher, Depends(_runner)],
-    task_status: str | SkipJsonSchema[None] = None,
+    task_status: TaskStatusEnum | None = None,
 ) -> TasksListResponse:
     """
     Retrieve tasks based on their status.
-    The status of a newly created task is 'unstarted'.
+    The status of a newly created task is PENDING.
     """
-    if task_status:
-        add_span_attributes({"status": task_status})
-        try:
-            desired_status = validate_task_status(task_status)
-        except ValueError as e:
-            raise HTTPException(
-                status_code=status.HTTP_400_BAD_REQUEST,
-                detail="Invalid status query parameter",
-            ) from e
-
-        tasks = runner.run(interface.get_tasks_by_status, desired_status)
-    else:
-        tasks = runner.run(interface.get_tasks)
+    tasks = runner.run(interface.get_tasks, task_status)
     return TasksListResponse(tasks=tasks)
 
 

src/blueapi/worker/task_worker.py

@@ -194,16 +194,6 @@ def cancel_active_task(
             add_span_attributes({"Task stopped": reason or default_reason})
         return self._current.task_id
 
-    @start_as_current_span(TRACER)
-    def get_tasks(self) -> list[TrackableTask]:
-        """
-        Return a list of all tasks on the worker,
-        any one of which can be triggered with begin_task.
-        Returns:
-            List[TrackableTask[T]]: List of task objects
-        """
-        return list(self._pending_tasks.values()) + list(self._completed_tasks.values())
-
     @start_as_current_span(TRACER, "task_id")
     def get_task_by_id(self, task_id: str) -> TrackableTask | None:
         """
@@ -217,12 +207,13 @@ def get_task_by_id(self, task_id: str) -> TrackableTask | None:
         """
         return self._pending_tasks.get(task_id, None) or self._completed_tasks[task_id]
 
-    @start_as_current_span(TRACER, "status")
-    def get_tasks_by_status(self, status: TaskStatusEnum) -> list[TrackableTask]:
+    @start_as_current_span(TRACER)
+    def get_tasks(self, status: TaskStatusEnum | None = None) -> list[TrackableTask]:
         """
         Retrieve a list of tasks based on their status.
         Args:
-           status TaskStatusEnum: The status to filter tasks by.
+           status Optional[TaskStatusEnum]: The status to filter tasks by.
+           If status is None return all tasks.
         Returns:
           list[TrackableTask]: A list of tasks that match the given status.
         """
@@ -236,6 +227,10 @@ def get_tasks_by_status(self, status: TaskStatusEnum) -> list[TrackableTask]:
             return [task for task in self._pending_tasks.values() if task.is_pending]
         elif status == TaskStatusEnum.COMPLETE:
             return list(self._completed_tasks.values())
+        elif status is None:
+            return list(self._pending_tasks.values()) + list(
+                self._completed_tasks.values()
+            )
         return []
 
     @start_as_current_span(TRACER)

tests/unit_tests/service/test_interface.py

@@ -262,28 +262,32 @@ def test_subscribers_removed_when_task_not_found(
     worker.worker_events.unsubscribe.assert_called_once_with(remove_token)
 
 
-@patch("blueapi.service.interface.TaskWorker.get_tasks_by_status")
-def test_get_tasks_by_status(get_tasks_by_status_mock: MagicMock):
-    pending_task1 = TrackableTask(task_id="0", task=Task(name="pending_task1"))
-    pending_task2 = TrackableTask(task_id="1", task=Task(name="pending_task2"))
-    running_task = TrackableTask(task_id="2", task=Task(name="running_task"))
+@patch("blueapi.service.interface.TaskWorker.get_tasks")
+def test_get_tasks(get_tasks_mock: MagicMock):
+    running_task = [TrackableTask(task_id="2", task=Task(name="running_task"))]
+    pending_task = [
+        TrackableTask(task_id="0", task=Task(name="pending_task1")),
+        TrackableTask(task_id="1", task=Task(name="pending_task2")),
+    ]
 
-    def mock_tasks_by_status(status: TaskStatusEnum) -> list[TrackableTask]:
+    def mock_tasks(
+        status: TaskStatusEnum | None = None,
+    ) -> list[TrackableTask]:
         if status == TaskStatusEnum.PENDING:
-            return [pending_task1, pending_task2]
+            return pending_task
         elif status == TaskStatusEnum.RUNNING:
-            return [running_task]
-        else:
+            return running_task
+        elif status == TaskStatusEnum.COMPLETE:
             return []
+        else:
+            return pending_task + running_task
 
-    get_tasks_by_status_mock.side_effect = mock_tasks_by_status
+    get_tasks_mock.side_effect = mock_tasks
 
-    assert interface.get_tasks_by_status(TaskStatusEnum.PENDING) == [
-        pending_task1,
-        pending_task2,
-    ]
-    assert interface.get_tasks_by_status(TaskStatusEnum.RUNNING) == [running_task]
-    assert interface.get_tasks_by_status(TaskStatusEnum.COMPLETE) == []
+    assert interface.get_tasks(TaskStatusEnum.PENDING) == pending_task
+    assert interface.get_tasks(TaskStatusEnum.RUNNING) == running_task
+    assert interface.get_tasks(TaskStatusEnum.COMPLETE) == []
+    assert interface.get_tasks() == pending_task + running_task
 
 
 @patch("blueapi.service.interface.BlueskyContext.numtracker")
@@ -334,18 +338,6 @@ def test_cancel_active_task(cancel_active_task_mock: MagicMock):
     cancel_active_task_mock.assert_called_once_with(fail, reason)
 
 
-@patch("blueapi.service.interface.TaskWorker.get_tasks")
-def test_get_tasks(get_tasks_mock: MagicMock):
-    tasks = [
-        TrackableTask(task_id="0", task=Task(name="0")),
-        TrackableTask(task_id="1", task=Task(name="1")),
-        TrackableTask(task_id="2", task=Task(name="2")),
-    ]
-    get_tasks_mock.return_value = tasks
-
-    assert interface.get_tasks() == tasks
-
-
 @pytest.mark.parametrize("tiled_enabled", [True, False])
 @patch("blueapi.service.interface.context")
 @patch("blueapi.service.interface.config")

tests/unit_tests/service/test_rest_api.py

@@ -411,9 +411,9 @@ def test_get_tasks_by_status(mock_runner: Mock, client: TestClient) -> None:
     }
 
 
-def test_get_tasks_by_status_invalid(client: TestClient) -> None:
+def test_get_tasks_by_invalid_status(client: TestClient) -> None:
     response = client.get("/tasks", params={"task_status": "AN_INVALID_STATUS"})
-    assert response.status_code == status.HTTP_400_BAD_REQUEST
+    assert response.status_code == status.HTTP_422_UNPROCESSABLE_CONTENT
 
 
 def test_delete_submitted_task(mock_runner: Mock, client: TestClient) -> None:

tests/unit_tests/worker/test_task_worker.py

@@ -628,7 +628,7 @@ def callback(unused: Future[list[Any]], stream=stream, sub=sub):
         (TaskStatusEnum.COMPLETE, ["task3"]),
     ],
 )
-def test_get_tasks_by_status(worker: TaskWorker, status, expected_task_ids):
+def test_get_tasks(worker: TaskWorker, status, expected_task_ids):
     worker._pending_tasks = {
         "task1": TrackableTask(
             task_id="task1",
@@ -658,7 +658,7 @@ def test_get_tasks_by_status(worker: TaskWorker, status, expected_task_ids):
         ),
     }
 
-    result = worker.get_tasks_by_status(status)
+    result = worker.get_tasks(status)
     result_ids = [task.task_id for task in result]
 
     assert result_ids == expected_task_ids