feat: add flaky test diagnostic script and fix flaky test

samples/snippets/conftest.py

@@ -180,9 +180,11 @@ def teardown_training_pipeline(shared_state, pipeline_client):
     )
 
     # Waiting for training pipeline to be in CANCELLED state
+    timeout = shared_state["cancel_batch_prediction_job_timeout"]
     helpers.wait_for_job_state(
         get_job_method=pipeline_client.get_training_pipeline,
         name=shared_state["training_pipeline_name"],
+        timeout=timeout,
     )
 
     # Delete the training pipeline

samples/snippets/helpers.py

@@ -1,5 +1,7 @@
+import collections
 import re
 import time
+from timeit import default_timer as timer
 
 from typing import Callable
 
@@ -63,3 +65,21 @@ def wait_for_job_state(
         "\nTry increasing the timeout in sample test"
         f"\nLast recorded state: {response.state}"
     )
+
+
+def flaky_test_diagnostic(file_name, test_name, N=20):
+
+    import pytest
+
+    timing_dict = collections.defaultdict(list)
+    for ri in range(N):
+        start = timer()
+        result = pytest.main(['-s', f'{file_name}::{test_name}'])
+        end = timer()
+        delta = end-start
+        if result == pytest.ExitCode.OK:
+            timing_dict['SUCCESS'].append(delta)
+        else:
+            timing_dict['FAILURE'].append(delta)
+
+    return timing_dict

samples/snippets/pipeline_service/create_training_pipeline_tabular_regression_sample_test.py

@@ -34,6 +34,9 @@ def teardown(teardown_training_pipeline):
 
 def test_ucaip_generated_create_training_pipeline_sample(capsys, shared_state):
 
+    # The return of the cancellation can be flaky; max of 20 runs was 215 sec.
+    shared_state["cancel_batch_prediction_job_timeout"] = 300
+
     create_training_pipeline_tabular_regression_sample.create_training_pipeline_tabular_regression_sample(
         project=PROJECT_ID,
         display_name=DISPLAY_NAME,

scripts/run_flaky_test_diagnostic.py

@@ -0,0 +1,22 @@
+import os
+import pathlib
+import sys
+
+# Run test in snippets directory:
+dir_of_curr_file = os.path.dirname(__file__)
+helper_filepath = pathlib.Path(dir_of_curr_file).parent / 'samples' / 'snippets'
+sys.path.append(helper_filepath.resolve().as_posix())
+os.chdir(helper_filepath.resolve())
+from helpers import flaky_test_diagnostic
+
+# Settings:
+file_name = 'pipeline_service/create_training_pipeline_tabular_regression_sample_test.py'
+test_name = 'test_ucaip_generated_create_training_pipeline_sample'
+timing_dict = flaky_test_diagnostic(file_name, test_name, N=1)
+
+for key, delta_list in timing_dict.items():
+    mean_time = sum(delta_list)/len(delta_list)
+    max_time = max(delta_list)
+    min_time = min(delta_list)
+    report_string = f'Result: {key}, mean={mean_time:3.2f}, min={min_time:3.2f}, max={max_time:3.2f}, count={len(delta_list)}'
+    print(report_string)