From a5a003b831a0578411995f26ae9938f8fccb8a31 Mon Sep 17 00:00:00 2001
From: Andrew Ferlitsch <aferlitsch@google.com>
Date: Tue, 24 Oct 2023 06:57:46 -0700
Subject: [PATCH] Harden 2 (#2414)

* feat: run-first argument

* feat: run-first code
---
 .cloud-build/execute_changed_notebooks_cli.py | 27 ++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/.cloud-build/execute_changed_notebooks_cli.py b/.cloud-build/execute_changed_notebooks_cli.py
index 55cd75410..3f1bf2a10 100755
--- a/.cloud-build/execute_changed_notebooks_cli.py
+++ b/.cloud-build/execute_changed_notebooks_cli.py
@@ -37,7 +37,7 @@ parser = argparse.ArgumentParser(description="Run changed notebooks.")
 parser.add_argument(
     "--test_paths_file",
     type=pathlib.Path,
-    help="The path to the file that has newline-limited folders of notebooks that should be tested.",
+    help="The path to the file that has newline-delimited folders of notebooks that should be tested.",
     required=True,
 )
 parser.add_argument(
@@ -128,6 +128,13 @@ parser.add_argument(
     default=10,
     required=False,
 )
+parser.add_argument(
+    "--run_first_file",
+    type=pathlib.Path,
+    help="The path to the file that has newline-delimited of notebooks to run in the first batch",
+    default=None,
+    required=False,
+)
 parser.add_argument(
     "--dry_run",
     type=str2bool,
@@ -158,6 +165,24 @@ else:
 
     notebooks = [changed_notebook for changed_notebook in changed_notebooks if execute_changed_notebooks_helper.select_notebook(changed_notebook, accumulative_results, args.test_percent)]
 
+run_first = []
+if args.run_first_file:
+    if not os.path.isfile(args.run_first_file):
+        print("Error: file does not exist", args.run_first_file)
+    else:
+        with open(args.run_first_file, 'r') as csvfile:
+            reader = csv.reader(csvfile)
+            for row in reader:
+                notebook = row[0]
+                run_first.append(notebook)
+
+    for notebook in run_first:
+        if notebook in notebooks:
+            # remove from existing list
+            notebooks.remove(notebook)
+            # add back to the front of the list
+            notebooks.insert(0, notebook)
+
 if args.dry_run:
     print("Dry run ...\n")
     for notebook in notebooks: