Except RayTaskError additionally in test_future_cancellation when can…

…celling future
aai-institute · Oct 26, 2023 · 5d409d1 · 5d409d1
1 parent 7dac359
commit 5d409d1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/utils/test_parallel.py b/tests/utils/test_parallel.py
@@ -243,7 +243,7 @@ def test_future_cancellation(parallel_config):
 
     assert future.result() == 2
 
-    from ray.exceptions import TaskCancelledError
+    from ray.exceptions import RayTaskError, TaskCancelledError
 
     with init_executor(
         config=parallel_config, cancel_futures=CancellationPolicy.ALL
@@ -256,7 +256,7 @@ def test_future_cancellation(parallel_config):
 
     assert future._state == "FINISHED"
 
-    with pytest.raises(TaskCancelledError):
+    with pytest.raises((TaskCancelledError, RayTaskError)):
         future.result()
 
     assert time.monotonic() - start < 1