From 5d409d111f3bb4baa811811c3a00cd445ee20d9c Mon Sep 17 00:00:00 2001 From: Anes Benmerzoug Date: Thu, 26 Oct 2023 13:55:15 +0200 Subject: [PATCH] Except RayTaskError additionally in test_future_cancellation when cancelling future --- tests/utils/test_parallel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/utils/test_parallel.py b/tests/utils/test_parallel.py index 18c2d9a64..22b28f5c2 100644 --- 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