@@ -683,7 +683,9 @@ def fake_sidecar_fct(
683683 )
684684 assert published_computation_task [0 ].node_id in image_params .fake_tasks
685685 # creating a new future shows that it is not done????
686- assert not distributed .Future (published_computation_task [0 ].job_id ).done ()
686+ assert not distributed .Future (
687+ published_computation_task [0 ].job_id , client = dask_client .backend .client
688+ ).done ()
687689
688690 # as the task is published on the dask-scheduler when sending, it shall still be published on the dask scheduler
689691 list_of_persisted_datasets = await dask_client .backend .client .list_datasets () # type: ignore
@@ -706,7 +708,9 @@ def fake_sidecar_fct(
706708 assert isinstance (task_result , TaskOutputData )
707709 assert task_result .get ("some_output_key" ) == 123
708710 # try to create another future and this one is already done
709- assert distributed .Future (published_computation_task [0 ].job_id ).done ()
711+ assert distributed .Future (
712+ published_computation_task [0 ].job_id , client = dask_client .backend .client
713+ ).done ()
710714
711715
712716async def test_abort_computation_tasks (
@@ -1111,7 +1115,9 @@ def fake_remote_fct(
11111115
11121116 assert published_computation_task [0 ].node_id in cpu_image .fake_tasks
11131117 # let's get a dask future for the task here so dask will not remove the task from the scheduler at the end
1114- computation_future = distributed .Future (key = published_computation_task [0 ].job_id )
1118+ computation_future = distributed .Future (
1119+ key = published_computation_task [0 ].job_id , client = dask_client .backend .client
1120+ )
11151121 assert computation_future
11161122
11171123 await _assert_wait_for_task_status (
0 commit comments