creating a test for this situation

Author: sanderegg

services/director-v2/tests/unit/with_dbs/comp_scheduler/test_scheduler_dask.py

@@ -2273,3 +2273,46 @@ async def test_run_new_pipeline_called_twice_prevents_duplicate_runs(
         0,  # No new messages expected
         ComputationalPipelineStatusMessage.model_validate_json,
     )
+
+
+@pytest.mark.parametrize(
+    "exception_type",
+    [
+        ComputationalBackendTaskResultsNotReadyError,
+    ],
+)
+async def test_getting_task_result_raises_exception(
+    exception_type: Exception,
+    with_disabled_auto_scheduling: mock.Mock,
+    with_disabled_scheduler_publisher: mock.Mock,
+    mocked_dask_client: mock.MagicMock,
+    initialized_app: FastAPI,
+    scheduler_api: BaseCompScheduler,
+    sqlalchemy_async_engine: AsyncEngine,
+    running_project: RunningProject,
+    mocked_parse_output_data_fct: mock.Mock,
+):
+    # this tests the behavior of the scheduling when the dask client cannot retrieve
+    # the result of a task because of some communication error. In this case the task
+    # it should be retrieved again in the next iteration and not marked as failed
+    # immediately.
+    async def mocked_get_tasks_status(job_ids: list[str]) -> list[RunningState]:
+        return [RunningState.SUCCESS for j in job_ids]
+
+    mocked_dask_client.get_tasks_status.side_effect = mocked_get_tasks_status
+    mocked_dask_client.get_task_result.side_effect = exception_type
+
+    # calling apply should not raise, but log the error
+    assert running_project.project.prj_owner
+    await scheduler_api.apply(
+        user_id=running_project.project.prj_owner,
+        project_id=running_project.project.uuid,
+        iteration=1,
+    )
+
+    assert running_project.project.prj_owner
+    await scheduler_api.apply(
+        user_id=running_project.project.prj_owner,
+        project_id=running_project.project.uuid,
+        iteration=1,
+    )