fix: review comment

.cloud-build/execute_changed_notebooks_cli.py

@@ -17,6 +17,7 @@
 
 import argparse
 import pathlib
+import random
 
 import execute_changed_notebooks_helper
 
@@ -39,6 +40,13 @@ parser.add_argument(
     help="The path to the file that has newline-limited folders of notebooks that should be tested.",
     required=True,
 )
+parser.add_argument(
+    "--test_percent",
+    type=int,
+    help="The percent of notebooks to be tested (between 1 and 100).",
+    required=False,
+    default=100,
+)
 parser.add_argument(
     "--base_branch",
     help="The base git branch to diff against to find changed files.",
@@ -107,24 +115,40 @@ parser.add_argument(
     default=True,
     help="Should run notebooks in parallel.",
 )
+parser.add_argument(
+    "--dry_run",
+    type=str2bool,
+    default=False,
+    help="Dry run for testing - no execution",
+)
 
 args = parser.parse_args()
 
-notebooks = execute_changed_notebooks_helper.get_changed_notebooks(
+changed_notebooks = execute_changed_notebooks_helper.get_changed_notebooks(
     test_paths_file=args.test_paths_file,
     base_branch=args.base_branch,
 )
 
-execute_changed_notebooks_helper.process_and_execute_notebooks(
-    notebooks=notebooks,
-    container_uri=args.container_uri,
-    staging_bucket=args.staging_bucket,
-    artifacts_bucket=args.artifacts_bucket,
-    should_parallelize=args.should_parallelize,
-    timeout=args.timeout,
-    variable_project_id=args.variable_project_id,
-    variable_region=args.variable_region,
-    variable_service_account=args.variable_service_account,
-    variable_vpc_network=args.variable_vpc_network,
-    private_pool_id=args.private_pool_id,
+if args.test_percent == 100:
+    notebooks = changed_notebooks
+else:
+    notebooks = [changed_notebook for changed_notebook in changed_notebooks if random.randint(1, 100) < args.test_percent]
+
+if args.dry_run:
+    print("Dry run ...\n")
+    for notebook in notebooks:
+        print(f"Would execute: {notebook.path}")
+else:
+    execute_changed_notebooks_helper.process_and_execute_notebooks(
+        notebooks=notebooks,
+        container_uri=args.container_uri,
+        staging_bucket=args.staging_bucket,
+        artifacts_bucket=args.artifacts_bucket,
+        should_parallelize=args.should_parallelize,
+        timeout=args.timeout,
+        variable_project_id=args.variable_project_id,
+        variable_region=args.variable_region,
+        variable_service_account=args.variable_service_account,
+        variable_vpc_network=args.variable_vpc_network,
+        private_pool_id=args.private_pool_id,
 )