From a5e6b651761a6e34b2dbfd3096f74741bdd05698 Mon Sep 17 00:00:00 2001
From: Charles Engelke <engelke@google.com>
Date: Wed, 26 Apr 2023 15:35:00 -0700
Subject: [PATCH] fix: accept flaky behavior of service (#9817)

* fix: accept flaky behavior of service

* fix: fix import statement
---
 batch/requirements-test.txt  | 1 +
 batch/tests/test_basics.py   | 4 ++++
 batch/tests/test_bucket.py   | 2 ++
 batch/tests/test_template.py | 3 +++
 4 files changed, 10 insertions(+)

diff --git a/batch/requirements-test.txt b/batch/requirements-test.txt
index 16f4b5c93074..82aa25583f33 100644
--- a/batch/requirements-test.txt
+++ b/batch/requirements-test.txt
@@ -2,3 +2,4 @@ pytest==7.2.0
 google-cloud-compute==1.11.0
 google-cloud-resource-manager==1.10.0
 google-cloud-storage==2.8.0
+flaky==3.7.0
\ No newline at end of file
diff --git a/batch/tests/test_basics.py b/batch/tests/test_basics.py
index 829a92729062..0998a63c76cb 100644
--- a/batch/tests/test_basics.py
+++ b/batch/tests/test_basics.py
@@ -15,6 +15,8 @@
 from typing import Callable
 import uuid
 
+from flaky import flaky
+
 import google.auth
 from google.cloud import batch_v1
 import pytest
@@ -90,11 +92,13 @@ def _check_logs(job, capsys):
     assert all("Hello world!" in log_msg for log_msg in output)
 
 
+@flaky(max_runs=3, min_passes=1)
 def test_script_job(job_name, capsys):
     job = create_script_job(PROJECT, REGION, job_name)
     _test_body(job, additional_test=lambda: _check_logs(job, capsys))
 
 
+@flaky(max_runs=3, min_passes=1)
 def test_container_job(job_name):
     job = create_container_job(PROJECT, REGION, job_name)
     _test_body(job, additional_test=lambda: _check_tasks(job_name))
diff --git a/batch/tests/test_bucket.py b/batch/tests/test_bucket.py
index ad8a347fbba0..b9eb3beff98f 100644
--- a/batch/tests/test_bucket.py
+++ b/batch/tests/test_bucket.py
@@ -13,6 +13,7 @@
 #  limitations under the License.
 import uuid
 
+from flaky import flaky
 
 import google.auth
 from google.cloud import batch_v1
@@ -65,6 +66,7 @@ def _test_bucket_content(test_bucket):
         assert content == file_content_template.format(task_number=i)
 
 
+@flaky(max_runs=3, min_passes=1)
 def test_bucket_job(job_name, test_bucket):
     job = create_script_job_with_bucket(PROJECT, REGION, job_name, test_bucket)
     _test_body(job, lambda: _test_bucket_content(test_bucket))
diff --git a/batch/tests/test_template.py b/batch/tests/test_template.py
index 572811731e0a..f7580e2f05d4 100644
--- a/batch/tests/test_template.py
+++ b/batch/tests/test_template.py
@@ -14,6 +14,8 @@
 
 import uuid
 
+from flaky import flaky
+
 import google.auth
 from google.cloud import batch_v1
 from google.cloud import compute_v1
@@ -109,6 +111,7 @@ def instance_template():
     operation_client.wait(project=PROJECT, operation=op.name)
 
 
+@flaky(max_runs=3, min_passes=1)
 def test_template_job(job_name, instance_template):
     job = create_script_job_with_template(PROJECT, REGION, job_name, instance_template.self_link)
     _test_body(job)