chore: import upstream snapshot with attribution

python/ray/autoscaler/v2/tests/test_ray_installer.py

@@ -0,0 +1,68 @@
+# coding: utf-8
+import os
+import sys
+import unittest
+
+import pytest  # noqa
+
+from ray._private.test_utils import load_test_config
+from ray.autoscaler.tags import TAG_RAY_NODE_KIND
+from ray.autoscaler.v2.instance_manager.config import AutoscalingConfig
+from ray.autoscaler.v2.instance_manager.ray_installer import RayInstaller
+from ray.core.generated.instance_manager_pb2 import Instance
+from ray.tests.autoscaler_test_utils import MockProcessRunner, MockProvider
+
+
+class RayInstallerTest(unittest.TestCase):
+    def setUp(self):
+        self.base_provider = MockProvider()
+        self.config = AutoscalingConfig(load_test_config("test_ray_complex.yaml"))
+        self.runner = MockProcessRunner()
+        self.ray_installer = RayInstaller(self.base_provider, self.config, self.runner)
+
+    def test_install_succeeded(self):
+        self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
+        self.runner.respond_to_call("json .Config.Env", ["[]" for i in range(1)])
+
+        self.ray_installer.install_ray(
+            Instance(
+                instance_id="0", instance_type="worker_nodes1", cloud_instance_id="0"
+            ),
+            head_node_ip="1.2.3.4",
+        )
+
+    def test_install_failed(self):
+        # creation failed because no such node.
+        with self.assertRaisesRegex(KeyError, "0"):
+            assert not self.ray_installer.install_ray(
+                Instance(
+                    instance_id="0",
+                    instance_type="worker_nodes1",
+                    cloud_instance_id="0",
+                ),
+                head_node_ip="1.2.3.4",
+            )
+
+        self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
+        self.runner.fail_cmds = [
+            "echo"  # this is the command used in the test_ray_complex.yaml
+        ]
+        self.runner.respond_to_call("json .Config.Env", ["[]" for i in range(1)])
+
+        # creation failed because setup command failed.
+        with self.assertRaisesRegex(Exception, "unexpected status"):
+            self.ray_installer.install_ray(
+                Instance(
+                    instance_id="0",
+                    instance_type="worker_nodes1",
+                    cloud_instance_id="0",
+                ),
+                head_node_ip="1.2.3.4",
+            )
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))