Refactor pytest fixtures for ray core (#4390)

python/ray/tests/test_multi_node.py

@@ -12,26 +12,8 @@ from ray.tests.utils import (run_and_get_output, run_string_as_driver,
                              run_string_as_driver_nonblocking)
 
 
-@pytest.fixture
-def ray_start_head():
-    out = run_and_get_output(["ray", "start", "--head", "--num-cpus=2"])
-    # Get the redis address from the output.
-    redis_substring_prefix = "redis_address=\""
-    redis_address_location = (
-        out.find(redis_substring_prefix) + len(redis_substring_prefix))
-    redis_address = out[redis_address_location:]
-    redis_address = redis_address.split("\"")[0]
-
-    yield redis_address
-
-    # Disconnect from the Ray cluster.
-    ray.shutdown()
-    # Kill the Ray cluster.
-    subprocess.Popen(["ray", "stop"]).wait()
-
-
-def test_error_isolation(ray_start_head):
-    redis_address = ray_start_head
+def test_error_isolation(call_ray_start):
+    redis_address = call_ray_start
     # Connect a driver to the Ray cluster.
     ray.init(redis_address=redis_address)
 
@@ -99,10 +81,10 @@ print("success")
     assert error_string1 in ray.error_info()[0]["message"]
 
 
-def test_remote_function_isolation(ray_start_head):
+def test_remote_function_isolation(call_ray_start):
     # This test will run multiple remote functions with the same names in
     # two different drivers. Connect a driver to the Ray cluster.
-    redis_address = ray_start_head
+    redis_address = call_ray_start
 
     ray.init(redis_address=redis_address)
 
@@ -142,10 +124,10 @@ print("success")
     assert "success" in out
 
 
-def test_driver_exiting_quickly(ray_start_head):
+def test_driver_exiting_quickly(call_ray_start):
     # This test will create some drivers that submit some tasks and then
     # exit without waiting for the tasks to complete.
-    redis_address = ray_start_head
+    redis_address = call_ray_start
 
     ray.init(redis_address=redis_address)
 
@@ -218,25 +200,11 @@ ray.get([a.log.remote(), f.remote()])
         assert out.count(log_message) == 4
 
 
-@pytest.fixture
-def ray_start_head_with_resources():
-    out = run_and_get_output(
-        ["ray", "start", "--head", "--num-cpus=1", "--num-gpus=1"])
-    # Get the redis address from the output.
-    redis_substring_prefix = "redis_address=\""
-    redis_address_location = (
-        out.find(redis_substring_prefix) + len(redis_substring_prefix))
-    redis_address = out[redis_address_location:]
-    redis_address = redis_address.split("\"")[0]
-
-    yield redis_address
-
-    # Kill the Ray cluster.
-    subprocess.Popen(["ray", "stop"]).wait()
-
-
-def test_drivers_release_resources(ray_start_head_with_resources):
-    redis_address = ray_start_head_with_resources
+@pytest.mark.parametrize(
+    "call_ray_start", ["ray start --head --num-cpus=1 --num-gpus=1"],
+    indirect=True)
+def test_drivers_release_resources(call_ray_start):
+    redis_address = call_ray_start
 
     # Define a driver that creates an actor and exits.
     driver_script1 = """
@@ -359,23 +327,13 @@ def test_calling_start_ray_head():
     subprocess.Popen(["ray", "stop"]).wait()
 
 
-@pytest.fixture
-def ray_start_head_local():
-    # Start the Ray processes on this machine.
-    run_and_get_output([
-        "ray", "start", "--head", "--node-ip-address=localhost",
-        "--redis-port=6379"
-    ])
-
-    yield None
-
-    # Disconnect from the Ray cluster.
-    ray.shutdown()
-    # Kill the Ray cluster.
-    subprocess.Popen(["ray", "stop"]).wait()
-
-
-def test_using_hostnames(ray_start_head_local):
+@pytest.mark.parametrize(
+    "call_ray_start", [
+        "ray start --head --num-cpus=1 " +
+        "--node-ip-address=localhost --redis-port=6379"
+    ],
+    indirect=True)
+def test_using_hostnames(call_ray_start):
     ray.init(node_ip_address="localhost", redis_address="localhost:6379")
 
     @ray.remote
@@ -385,15 +343,6 @@ def test_using_hostnames(ray_start_head_local):
     assert ray.get(f.remote()) == 1
 
 
-@pytest.fixture
-def ray_start_regular():
-    # Start the Ray processes.
-    address_info = ray.init(num_cpus=1)
-    yield address_info
-    # The code after the yield will run as teardown code.
-    ray.shutdown()
-
-
 def test_connecting_in_local_case(ray_start_regular):
     address_info = ray_start_regular
 
@@ -453,10 +402,10 @@ print("success")
         assert "success" in out
 
 
-def test_driver_exiting_when_worker_blocked(ray_start_head):
+def test_driver_exiting_when_worker_blocked(call_ray_start):
     # This test will create some drivers that submit some tasks and then
     # exit without waiting for the tasks to complete.
-    redis_address = ray_start_head
+    redis_address = call_ray_start
 
     ray.init(redis_address=redis_address)