From bdf0c009893c1e153543766ce3941ce99084afa2 Mon Sep 17 00:00:00 2001
From: Ian Rodney <ian.rodney@gmail.com>
Date: Wed, 27 Jan 2021 15:33:33 -0800
Subject: [PATCH] Revert "Revert "[CLI] Fix Ray Status with ENV Variable set
 (#13707) (#13726)

---
 python/ray/_private/services.py               |  2 +-
 python/ray/tests/test_cli.py                  | 28 +++++++++++++++++++
 .../test_cli_patterns/test_ray_status.txt     | 14 ++++++++++
 3 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 python/ray/tests/test_cli_patterns/test_ray_status.txt

diff --git a/python/ray/_private/services.py b/python/ray/_private/services.py
index d0eafc969..1c4c6497d 100644
--- a/python/ray/_private/services.py
+++ b/python/ray/_private/services.py
@@ -216,7 +216,7 @@ def get_ray_address_to_use_or_die():
         A string to pass into `ray.init(address=...)`
     """
     if "RAY_ADDRESS" in os.environ:
-        return "auto"  # Avoid conflict with RAY_ADDRESS env var
+        return os.environ.get("RAY_ADDRESS")
 
     return find_redis_address_or_die()
 
diff --git a/python/ray/tests/test_cli.py b/python/ray/tests/test_cli.py
index 57bf61419..f5628701f 100644
--- a/python/ray/tests/test_cli.py
+++ b/python/ray/tests/test_cli.py
@@ -37,6 +37,7 @@ from testfixtures.popen import MockPopen, PopenBehaviour
 
 import ray.autoscaler._private.aws.config as aws_config
 import ray.scripts.scripts as scripts
+from ray.test_utils import wait_for_condition
 
 boto3_list = [{
     "InstanceType": "t1.micro",
@@ -415,5 +416,32 @@ def test_ray_submit(configure_lang, configure_aws, _unlink_test_ssh_key):
             _check_output_via_pattern("test_ray_submit.txt", result)
 
 
+def test_ray_status():
+    import ray
+    address = ray.init().get("redis_address")
+    runner = CliRunner()
+
+    def output_ready():
+        result = runner.invoke(scripts.status)
+        result.stdout
+        return not result.exception and "memory" in result.output
+
+    wait_for_condition(output_ready)
+
+    result = runner.invoke(scripts.status, [])
+    _check_output_via_pattern("test_ray_status.txt", result)
+
+    result_arg = runner.invoke(scripts.status, ["--address", address])
+    _check_output_via_pattern("test_ray_status.txt", result_arg)
+
+    # Try to check status with RAY_ADDRESS set
+    os.environ["RAY_ADDRESS"] = address
+    result_env = runner.invoke(scripts.status)
+    _check_output_via_pattern("test_ray_status.txt", result_env)
+
+    result_env_arg = runner.invoke(scripts.status, ["--address", address])
+    _check_output_via_pattern("test_ray_status.txt", result_env_arg)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))
diff --git a/python/ray/tests/test_cli_patterns/test_ray_status.txt b/python/ray/tests/test_cli_patterns/test_ray_status.txt
new file mode 100644
index 000000000..f903c6d62
--- /dev/null
+++ b/python/ray/tests/test_cli_patterns/test_ray_status.txt
@@ -0,0 +1,14 @@
+======== Cluster status: .+
+Node status
+------------------------------------------------------------
+ 1 node\(s\) with resources: .+
+
+Resources
+------------------------------------------------------------
+Usage:
+ 0.+
+ 0.+
+ 0.+
+
+Demands:
+ \(no resource demands\)