added backoff to one test

datalabeling/create_annotation_spec_set_test.py

@@ -16,30 +16,36 @@
 
 import os
 
-from google.api_core.client_options import ClientOptions
-from google.cloud import datalabeling_v1beta1 as datalabeling
+import backoff
+from google.api_core.exceptions import DeadlineExceeded
 import pytest
 
 import create_annotation_spec_set
+import testing_lib
+
 
 PROJECT_ID = os.getenv('GCLOUD_PROJECT')
 
 
-@pytest.mark.flaky(max_runs=3)
+@pytest.fixture(scope='module')
+def cleaner():
+    resource_names = []
+
+    yield resource_names
+
+    for resource_name in resource_names:
+        testing_lib.delete_annotation_spec_set(resource_name)
+
+
 def test_create_annotation_spec_set(capsys):
-    response = create_annotation_spec_set.create_annotation_spec_set(
-        PROJECT_ID)
-    out, _ = capsys.readouterr()
-    assert 'The annotation_spec_set resource name:' in out
 
-    # Delete the created annotation spec set.
-    annotation_spec_set_name = response.name
-    client = datalabeling.DataLabelingServiceClient()
+    @backoff.on_exception(backoff.expo, DeadlineExceeded, max_time=60)
+    def run_sample():
+        return create_annotation_spec_set.create_annotation_spec_set(PROJECT_ID)
 
-    # If provided, use a provided test endpoint - this will prevent tests on
-    # this snippet from triggering any action by a real human
-    if 'DATALABELING_ENDPOINT' in os.environ:
-        opts = ClientOptions(api_endpoint=os.getenv('DATALABELING_ENDPOINT'))
-        client = datalabeling.DataLabelingServiceClient(client_options=opts)
+    response = run_sample()
+    out, _ = capsys.readouterr()
+    assert 'The annotation_spec_set resource name:' in out
 
-    client.delete_annotation_spec_set(annotation_spec_set_name)
+    # For cleanup
+    cleaner.append(response.name)

datalabeling/requirements-test.txt

@@ -1,2 +1,3 @@
+backoff==1.10.0
 pytest==5.3.2
 flaky==3.6.1

datalabeling/testing_lib.py

@@ -0,0 +1,37 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import backoff
+from google.api_core.client_options import ClientOptions
+from google.api_core.exceptions import DeadlineExceeded
+from google.cloud import datalabeling_v1beta1 as datalabeling
+
+
+def create_client():
+    # If provided, use a provided test endpoint - this will prevent tests on
+    # this snippet from triggering any action by a real human
+    if 'DATALABELING_ENDPOINT' in os.environ:
+        opts = ClientOptions(api_endpoint=os.getenv('DATALABELING_ENDPOINT'))
+        client = datalabeling.DataLabelingServiceClient(client_options=opts)
+    else:
+        client = datalabeling.DataLabelingServiceClient()
+    return client
+
+
+@backoff.on_exception(backoff.expo, DeadlineExceeded, max_time=60)
+def delete_annotation_spec_set(name):
+    client = create_client()
+    client.delete_annotation_spec_set(name)