From 557da7044ffe7d9a85c8b8e28512c648659e5eee Mon Sep 17 00:00:00 2001
From: SangBin Cho <rkooo567@gmail.com>
Date: Mon, 6 Jul 2020 15:19:00 -0700
Subject: [PATCH] Fix flaky test that says `ray.init` is called twice. (#9234)

---
 python/ray/tests/test_actor_failures.py | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/python/ray/tests/test_actor_failures.py b/python/ray/tests/test_actor_failures.py
index 1d4fb5a08..fd012bd0f 100644
--- a/python/ray/tests/test_actor_failures.py
+++ b/python/ray/tests/test_actor_failures.py
@@ -90,6 +90,16 @@ def ray_checkpointable_actor_cls(request):
     return CheckpointableActor
 
 
+@pytest.fixture
+def ray_init_with_task_retry_delay():
+    address = ray.init(
+        _internal_config=json.dumps({
+            "task_retry_delay_ms": 100
+        }))
+    yield address
+    ray.shutdown()
+
+
 @pytest.mark.parametrize(
     "ray_start_regular", [{
         "object_store_memory": 150 * 1024 * 1024,
@@ -134,12 +144,8 @@ def test_actor_eviction(ray_start_regular):
     assert num_success > 0
 
 
-def test_actor_restart():
+def test_actor_restart(ray_init_with_task_retry_delay):
     """Test actor restart when actor process is killed."""
-    ray.init(
-        _internal_config=json.dumps({
-            "task_retry_delay_ms": 100,
-        }), )
 
     @ray.remote(max_restarts=1, max_task_retries=-1)
     class RestartableActor:
@@ -213,15 +219,10 @@ def test_actor_restart():
     # Check that the actor won't be restarted.
     with pytest.raises(ray.exceptions.RayActorError):
         ray.get(actor.increase.remote())
-    ray.shutdown()
 
 
-def test_actor_restart_with_retry():
+def test_actor_restart_with_retry(ray_init_with_task_retry_delay):
     """Test actor restart when actor process is killed."""
-    ray.init(
-        _internal_config=json.dumps({
-            "task_retry_delay_ms": 100,
-        }), )
 
     @ray.remote(max_restarts=1, max_task_retries=-1)
     class RestartableActor:
@@ -274,7 +275,6 @@ def test_actor_restart_with_retry():
     # Check that the actor won't be restarted.
     with pytest.raises(ray.exceptions.RayActorError):
         ray.get(actor.increase.remote())
-    ray.shutdown()
 
 
 def test_actor_restart_on_node_failure(ray_start_cluster):