fixed usage of random

Author: sanderegg

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

@@ -27,7 +27,10 @@
     assert_comp_runs_empty,
     assert_comp_tasks_and_comp_run_snapshot_tasks,
 )
-from dask_task_models_library.container_tasks.errors import TaskCancelledError
+from dask_task_models_library.container_tasks.errors import (
+    ServiceRuntimeError,
+    TaskCancelledError,
+)
 from dask_task_models_library.container_tasks.events import TaskProgressEvent
 from dask_task_models_library.container_tasks.io import TaskOutputData
 from dask_task_models_library.container_tasks.protocol import TaskOwner
@@ -838,7 +841,13 @@ async def _return_2nd_task_failed(job_ids: list[str]) -> list[RunningState]:
         ]
 
     mocked_dask_client.get_tasks_status.side_effect = _return_2nd_task_failed
-    mocked_dask_client.get_task_result.side_effect = None
+    mocked_dask_client.get_task_result.side_effect = ServiceRuntimeError(
+        service_key="simcore/services/dynamic/some-service",
+        service_version="1.0.0",
+        container_id="some-container-id",
+        exit_code=1,
+        service_logs="simulated error",
+    )
     await scheduler_api.apply(
         user_id=run_in_db.user_id,
         project_id=run_in_db.project_uuid,
@@ -2301,18 +2310,17 @@ async def mocked_get_tasks_status(job_ids: list[str]) -> list[RunningState]:
     computational_tasks = [
         t for t in running_project.tasks if t.node_class is NodeClass.COMPUTATIONAL
     ]
-    expected_timeouted_tasks = random.choices(  # noqa: S311
+    expected_timeouted_tasks = random.sample(
         computational_tasks, k=len(computational_tasks) - 1
     )
     successful_tasks = [
         t for t in computational_tasks if t not in expected_timeouted_tasks
     ]
 
     async def mocked_get_task_result(job_id: str) -> TaskOutputData:
-        nonlocal expected_timeouted_tasks
-        if job_id in [t.job_id for t in expected_timeouted_tasks]:
-            raise ComputationalBackendTaskResultsNotReadyError(job_id=job_id)
-        return TaskOutputData.model_validate({"whatever_output": 123})
+        if job_id in [t.job_id for t in successful_tasks]:
+            return TaskOutputData.model_validate({"whatever_output": 123})
+        raise ComputationalBackendTaskResultsNotReadyError(job_id=job_id)
 
     mocked_dask_client.get_task_result.side_effect = mocked_get_task_result
     # calling apply should not raise
@@ -2346,6 +2354,7 @@ async def mocked_get_task_result(job_id: str) -> TaskOutputData:
         retrieval_times.append(
             error_dict["ctx"][_TASK_RETRIEVAL_ERROR_CONTEXT_TIME_KEY]
         )
+    assert len(retrieval_times) == len(expected_timeouted_tasks)
 
     await assert_comp_tasks_and_comp_run_snapshot_tasks(
         sqlalchemy_async_engine,
@@ -2374,14 +2383,7 @@ async def mocked_get_task_result(job_id: str) -> TaskOutputData:
             iteration=1,
         )
         assert mocked_dask_client.get_task_result.call_count == (
-            len(
-                [
-                    t
-                    for t in running_project.tasks
-                    if t.node_class is NodeClass.COMPUTATIONAL
-                ]
-            )
-            - 1
+            len(expected_timeouted_tasks)
         )
         mocked_dask_client.get_task_result.reset_mock()
         await asyncio.sleep(0.5)  # wait a bit to ensure the retry decorator has reset