[projects] Add named commands to sessions (#5525)

python/ray/tests/test_projects.py

@@ -87,10 +87,9 @@ def _chdir_and_back(d):
         os.chdir(old_dir)
 
 
-def test_session_start_default_project():
+def run_test_project(project_dir, command, args):
     # Run the CLI commands with patching
-    test_dir = os.path.join(TEST_DIR,
-                            "project_files/session-tests/project-pass")
+    test_dir = os.path.join(TEST_DIR, "project_files", project_dir)
     with _chdir_and_back(test_dir):
         runner = CliRunner()
         with patch.multiple(
@@ -99,11 +98,17 @@ def test_session_start_default_project():
                 rsync=DEFAULT,
                 exec_cluster=DEFAULT,
         ) as mock_calls:
-            result = runner.invoke(start, [])
-            assert result.exit_code == 0
+            result = runner.invoke(command, args)
+
+    return result, mock_calls, test_dir
+
+
+def test_session_start_default_project():
+    result, mock_calls, test_dir = run_test_project(
+        "session-tests/project-pass", start, [])
 
-    # Check we are calling autoscaler correctly
     loaded_project = ray.projects.load_project(test_dir)
+    assert result.exit_code == 0
 
     # Part 1/3: Cluster Launching Call
     create_or_update_cluster_call = mock_calls["create_or_update_cluster"]
@@ -150,39 +155,20 @@ def test_session_start_default_project():
 
 
 def test_session_start_docker_fail():
-    # Run the CLI commands with patching
-    test_dir = os.path.join(TEST_DIR,
-                            "project_files/session-tests/with-docker-fail")
-    with _chdir_and_back(test_dir):
-        runner = CliRunner()
-        with patch.multiple(
-                "ray.projects.scripts",
-                create_or_update_cluster=DEFAULT,
-                rsync=DEFAULT,
-                exec_cluster=DEFAULT,
-        ) as _:
-            result = runner.invoke(start, [])
-            assert result.exit_code == 1
-            assert ("Docker support in session is currently "
-                    "not implemented") in result.output
+    result, _, _ = run_test_project("session-tests/with-docker-fail", start,
+                                    [])
+
+    assert result.exit_code == 1
+    assert ("Docker support in session is currently "
+            "not implemented") in result.output
 
 
 def test_session_git_repo_cloned():
-    # Run the CLI commands with patching
-    test_dir = os.path.join(TEST_DIR,
-                            "project_files/session-tests/git-repo-pass")
-    with _chdir_and_back(test_dir):
-        runner = CliRunner()
-        with patch.multiple(
-                "ray.projects.scripts",
-                create_or_update_cluster=DEFAULT,
-                rsync=DEFAULT,
-                exec_cluster=DEFAULT,
-        ) as mock_calls:
-            result = runner.invoke(start, [])
-            assert result.exit_code == 0
+    result, mock_calls, test_dir = run_test_project(
+        "session-tests/git-repo-pass", start, [])
 
     loaded_project = ray.projects.load_project(test_dir)
+    assert result.exit_code == 0
 
     exec_cluster_call = mock_calls["exec_cluster"]
     commands_executed = []
@@ -197,19 +183,27 @@ def test_session_git_repo_cloned():
 
 
 def test_session_invalid_config_errored():
-    # Run the CLI commands with patching
-    test_dir = os.path.join(TEST_DIR,
-                            "project_files/session-tests/invalid-config-fail")
-    with _chdir_and_back(test_dir):
-        runner = CliRunner()
-        with patch.multiple(
-                "ray.projects.scripts",
-                create_or_update_cluster=DEFAULT,
-                rsync=DEFAULT,
-                exec_cluster=DEFAULT,
-        ) as _:
-            result = runner.invoke(start, [])
-            assert result.exit_code == 1
-            assert "validation failed" in result.output
-            # check that we are displaying actional error message
-            assert "ray project validate" in result.output
+    result, _, _ = run_test_project("session-tests/invalid-config-fail", start,
+                                    [])
+
+    assert result.exit_code == 1
+    assert "validation failed" in result.output
+    # check that we are displaying actional error message
+    assert "ray project validate" in result.output
+
+
+def test_session_create_command():
+    result, mock_calls, test_dir = run_test_project(
+        "session-tests/commands-test", start,
+        ["first", "--a", "1", "--b", "2"])
+
+    # Verify the project can be loaded.
+    ray.projects.load_project(test_dir)
+    assert result.exit_code == 0
+
+    exec_cluster_call = mock_calls["exec_cluster"]
+    found_command = False
+    for _, kwargs in exec_cluster_call.call_args_list:
+        if "Starting ray job with 1 and 2" in kwargs["cmd"]:
+            found_command = True
+    assert found_command