From f730d6b9de58efedf736265e99cd7ffcb7cb75b7 Mon Sep 17 00:00:00 2001
From: Ivan Cheung <ivans.mailbox@gmail.com>
Date: Mon, 6 Jun 2022 10:30:31 -0400
Subject: [PATCH] Added ability to test a single notebook (#608)

* Added ability to test a single notebook

* Added output_url to table

* Removed ML Ops notebooks
---
 .../execute_changed_notebooks_helper.py       | 95 +++++++++++--------
 .cloud-build/test_folders.txt                 |  1 -
 2 files changed, 55 insertions(+), 41 deletions(-)

diff --git a/.cloud-build/execute_changed_notebooks_helper.py b/.cloud-build/execute_changed_notebooks_helper.py
index fc95461ba..a0ab91840 100755
--- a/.cloud-build/execute_changed_notebooks_helper.py
+++ b/.cloud-build/execute_changed_notebooks_helper.py
@@ -24,6 +24,7 @@ import re
 import subprocess
 from typing import List, Optional
 
+import execute_notebook_helper
 import execute_notebook_remote
 import nbformat
 from google.cloud.devtools.cloudbuild_v1.types import BuildOperationMetadata
@@ -287,14 +288,15 @@ def process_and_execute_notebooks(
         timeout (str):
             Required. Timeout string according to https://cloud.google.com/build/docs/build-config-file-schema#timeout.
     """
-    notebook_execution_results: List[NotebookExecutionResult] = []
 
     # Calculate deadline
     deadline = datetime.datetime.now() + datetime.timedelta(
         seconds=max(timeout - WORKER_TIMEOUT_BUFFER_IN_SECONDS, 0)
     )
 
-    if len(notebooks) > 0:
+    if len(notebooks) > 1:
+        notebook_execution_results: List[NotebookExecutionResult] = []
+
         print(f"Found {len(notebooks)} modified notebooks: {notebooks}")
 
         if should_parallelize and len(notebooks) > 1:
@@ -333,43 +335,56 @@ def process_and_execute_notebooks(
                 )
                 for notebook in notebooks
             ]
+
+        print("\n=== RESULTS ===\n")
+
+        results_sorted = sorted(
+            notebook_execution_results,
+            key=lambda result: result.is_pass,
+            reverse=True,
+        )
+
+        # Print results
+        print(
+            tabulate(
+                [
+                    [
+                        result.name,
+                        "PASSED" if result.is_pass else "FAILED",
+                        format_timedelta(result.duration),
+                        result.log_url,
+                        result.output_uri,
+                    ]
+                    for result in results_sorted
+                ],
+                headers=["build_tag", "status", "duration", "log_url", "output_url"],
+            )
+        )
+
+        print("\n=== END RESULTS===\n")
+
+        total_notebook_duration = functools.reduce(
+            operator.add,
+            [datetime.timedelta(seconds=0)]
+            + [result.duration for result in results_sorted],
+        )
+
+        print(
+            f"Cumulative notebook duration: {format_timedelta(total_notebook_duration)}"
+        )
+
+        # Raise error if any notebooks failed
+        if not all([result.is_pass for result in results_sorted]):
+            raise RuntimeError("Notebook failures detected. See logs for details")
+
+    elif len(notebooks) == 1:
+        notebook = notebooks[0]
+        execute_notebook_helper.execute_notebook(
+            notebook_source=notebook,
+            output_file_or_uri="/".join(
+                [artifacts_bucket, pathlib.Path(notebook).name]
+            ),
+            should_log_output=True,
+        )
     else:
         print("No notebooks modified in this pull request.")
-
-    print("\n=== RESULTS ===\n")
-
-    results_sorted = sorted(
-        notebook_execution_results,
-        key=lambda result: result.is_pass,
-        reverse=True,
-    )
-
-    # Print results
-    print(
-        tabulate(
-            [
-                [
-                    result.name,
-                    "PASSED" if result.is_pass else "FAILED",
-                    format_timedelta(result.duration),
-                    result.log_url,
-                ]
-                for result in results_sorted
-            ],
-            headers=["build_tag", "status", "duration", "log_url"],
-        )
-    )
-
-    print("\n=== END RESULTS===\n")
-
-    total_notebook_duration = functools.reduce(
-        operator.add,
-        [datetime.timedelta(seconds=0)]
-        + [result.duration for result in results_sorted],
-    )
-
-    print(f"Cumulative notebook duration: {format_timedelta(total_notebook_duration)}")
-
-    # Raise error if any notebooks failed
-    if not all([result.is_pass for result in results_sorted]):
-        raise RuntimeError("Notebook failures detected. See logs for details")
diff --git a/.cloud-build/test_folders.txt b/.cloud-build/test_folders.txt
index 69c98e768..502bc5d72 100644
--- a/.cloud-build/test_folders.txt
+++ b/.cloud-build/test_folders.txt
@@ -1,3 +1,2 @@
 notebooks/official
 notebooks/notebook_template.ipynb
-notebooks/community/ml_ops