fix: accept flaky behavior of service (GoogleCloudPlatform#9817)

* fix: accept flaky behavior of service

* fix: fix import statement

Author: engelke

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

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))

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))

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)