use backoff instead of flaky

GoogleCloudPlatform · Apr 22, 2020 · aaff0be · aaff0be
1 parent 9fa6ea4
commit aaff0be
Show file tree

Hide file tree

Showing 9 changed files with 243 additions and 210 deletions.
diff --git a/datalabeling/create_annotation_spec_set_test.py b/datalabeling/create_annotation_spec_set_test.py
@@ -39,7 +39,8 @@ def cleaner():
 
 def test_create_annotation_spec_set(cleaner, capsys):
 
-    @backoff.on_exception(backoff.expo, DeadlineExceeded, max_time=60)
+    @backoff.on_exception(
+        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
     def run_sample():
         return create_annotation_spec_set.create_annotation_spec_set(PROJECT_ID)
 

diff --git a/datalabeling/create_instruction_test.py b/datalabeling/create_instruction_test.py
@@ -16,35 +16,39 @@
 
 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_instruction
+import testing_lib
+
 
 PROJECT_ID = os.getenv('GCLOUD_PROJECT')
 INSTRUCTION_GCS_URI = ('gs://cloud-samples-data/datalabeling'
                        '/instruction/test.pdf')
 
 
-@pytest.mark.flaky(max_runs=3)
-def test_create_instruction(capsys):
-    result = create_instruction.create_instruction(
-        PROJECT_ID,
-        'IMAGE',
-        INSTRUCTION_GCS_URI
-    )
-    out, _ = capsys.readouterr()
-    assert 'The instruction resource name: ' in out
+@pytest.fixture(scope='module')
+def cleaner():
+    resource_names = []
+
+    yield resource_names
+
+    for resource_name in resource_names:
+        testing_lib.delete_instruction(resource_name)
 
-    # Delete the created instruction.
-    instruction_name = result.name
-    client = datalabeling.DataLabelingServiceClient()
 
-    # 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)
+def test_create_instruction(cleaner, capsys):
 
-    client.delete_instruction(instruction_name)
+    @backoff.on_exception(
+        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+    def run_sample():
+        return create_instruction.create_instruction(
+            PROJECT_ID, 'IMAGE', INSTRUCTION_GCS_URI)
+
+    instruction = run_sample()
+    cleaner.append(instruction.name)
+
+    out, _ = capsys.readouterr()
+    assert 'The instruction resource name: ' in out
diff --git a/datalabeling/import_data_test.py b/datalabeling/import_data_test.py
@@ -16,28 +16,48 @@
 
 import os
 
+import backoff
+from google.api_core.exceptions import DeadlineExceeded
 import pytest
 
 import import_data
 import manage_dataset
+import testing_lib
+
 
 PROJECT_ID = os.getenv('GCLOUD_PROJECT')
 INPUT_GCS_URI = 'gs://cloud-samples-data/datalabeling/image/image_dataset.csv'
 
 
 @pytest.fixture(scope='function')
 def dataset():
-    # create a temporary dataset
-    dataset = manage_dataset.create_dataset(PROJECT_ID)
+
+    @backoff.on_exception(
+        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+    def create_dataset():
+        # create a temporary dataset
+        return manage_dataset.create_dataset(PROJECT_ID)
+
+    dataset = create_dataset()
 
     yield dataset
 
-    # tear down
-    manage_dataset.delete_dataset(dataset.name)
+    @backoff.on_exception(
+        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+    def delete_dataset():
+        # tear down
+        manage_dataset.delete_dataset(dataset.name)
+
+    delete_dataset()
 
 
-@pytest.mark.flaky(max_runs=3)
 def test_import_data(capsys, dataset):
-    import_data.import_data(dataset.name, 'IMAGE', INPUT_GCS_URI)
+
+    @backoff.on_exception(
+        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+    def run_sample():
+        import_data.import_data(dataset.name, 'IMAGE', INPUT_GCS_URI)
+
+    run_sample()
     out, _ = capsys.readouterr()
     assert 'Dataset resource name: ' in out
diff --git a/datalabeling/label_image_test.py b/datalabeling/label_image_test.py
@@ -16,101 +16,77 @@
 
 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 create_instruction
-import import_data
 import label_image
-import manage_dataset
+import testing_lib
+
 
 PROJECT_ID = os.getenv('GCLOUD_PROJECT')
 INPUT_GCS_URI = 'gs://cloud-samples-data/datalabeling/image/image_dataset.csv'
+INSTRUCTION_GCS_URI = ('gs://cloud-samples-data/datalabeling'
+                       '/instruction/test.pdf')
 
 
 @pytest.fixture(scope='function')
 def dataset():
     # create a temporary dataset
-    dataset = manage_dataset.create_dataset(PROJECT_ID)
-
-    # import some data to it
-    import_data.import_data(dataset.name, 'IMAGE', INPUT_GCS_URI)
+    dataset = testing_lib.create_dataset(PROJECT_ID)
 
     yield dataset
 
     # tear down
-    manage_dataset.delete_dataset(dataset.name)
+    testing_lib.delete_dataset(dataset.name)
 
 
 @pytest.fixture(scope='function')
 def annotation_spec_set():
     # create a temporary annotation_spec_set
-    response = create_annotation_spec_set.create_annotation_spec_set(
-        PROJECT_ID)
+    response = testing_lib.create_annotation_spec_set(PROJECT_ID)
 
     yield response
 
-    # tear down
-    client = datalabeling.DataLabelingServiceClient()
-
-    # 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)
-
-    client.delete_annotation_spec_set(response.name)
+    testing_lib.delete_annotation_spec_set(response.name)
 
 
 @pytest.fixture(scope='function')
 def instruction():
     # create a temporary instruction
-    instruction = create_instruction.create_instruction(
-            PROJECT_ID, 'IMAGE',
-            'gs://cloud-samples-data/datalabeling/instruction/test.pdf')
+    instruction = testing_lib.create_instruction(
+        PROJECT_ID, 'IMAGE', INSTRUCTION_GCS_URI)
 
     yield instruction
 
     # tear down
-    client = datalabeling.DataLabelingServiceClient()
+    testing_lib.delete_instruction(instruction.name)
+
 
-    # 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)
+@pytest.fixture(scope='module')
+def cleaner():
+    resource_names = []
 
-    client.delete_instruction(instruction.name)
+    yield resource_names
+
+    for resource_name in resource_names:
+        testing_lib.cancel_operation(resource_name)
 
 
 # Passing in dataset as the last argument in test_label_image since it needs
 # to be deleted before the annotation_spec_set can be deleted.
-@pytest.mark.flaky(max_runs=3)
-def test_label_image(capsys, annotation_spec_set, instruction, dataset):
-
-    # Start labeling.
-    response = label_image.label_image(
-        dataset.name,
-        instruction.name,
-        annotation_spec_set.name
-    )
-    out, _ = capsys.readouterr()
-    assert 'Label_image operation name: ' in out
-    operation_name = response.operation.name
+def test_label_image(
+        capsys, annotation_spec_set, instruction, dataset, cleaner):
 
-    # Cancels the labeling operation.
-    response.cancel()
-    assert response.cancelled() is True
+    @backoff.on_exception(
+        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+    def run_sample():
+        # Start labeling.
+        return label_image.label_image(
+            dataset.name, instruction.name, annotation_spec_set.name)
 
-    client = datalabeling.DataLabelingServiceClient()
+    response = run_sample()
+    cleaner.append(response.operation.name)
 
-    # 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)
-
-    client.transport._operations_client.cancel_operation(
-            operation_name)
+    out, _ = capsys.readouterr()
+    assert 'Label_image operation name: ' in out
diff --git a/datalabeling/label_text_test.py b/datalabeling/label_text_test.py
@@ -16,101 +16,81 @@
 
 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 create_instruction
-import import_data
 import label_text
-import manage_dataset
+import testing_lib
 
 PROJECT_ID = os.getenv('GCLOUD_PROJECT')
 INPUT_GCS_URI = 'gs://cloud-samples-data/datalabeling/text/input.csv'
+INSTRUCTION_GCS_URI = ('gs://cloud-samples-data/datalabeling'
+                       '/instruction/test.pdf')
 
 
 @pytest.fixture(scope='function')
 def dataset():
     # create a temporary dataset
-    dataset = manage_dataset.create_dataset(PROJECT_ID)
+    dataset = testing_lib.create_dataset(PROJECT_ID)
 
-    # import some data to it
-    import_data.import_data(dataset.name, 'TEXT', INPUT_GCS_URI)
+    testing_lib.import_data(dataset.name, 'TEXT', INPUT_GCS_URI)
 
     yield dataset
 
     # tear down
-    manage_dataset.delete_dataset(dataset.name)
+    testing_lib.delete_dataset(dataset.name)
 
 
 @pytest.fixture(scope='function')
 def annotation_spec_set():
     # create a temporary annotation_spec_set
-    response = create_annotation_spec_set.create_annotation_spec_set(
-        PROJECT_ID)
+    response = testing_lib.create_annotation_spec_set(PROJECT_ID)
 
     yield response
 
-    # tear down
-    client = datalabeling.DataLabelingServiceClient()
-
-    # 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)
-
-    client.delete_annotation_spec_set(response.name)
+    testing_lib.delete_annotation_spec_set(response.name)
 
 
 @pytest.fixture(scope='function')
 def instruction():
     # create a temporary instruction
-    instruction = create_instruction.create_instruction(
-            PROJECT_ID, 'TEXT',
-            'gs://cloud-samples-data/datalabeling/instruction/test.pdf')
+    instruction = testing_lib.create_instruction(
+        PROJECT_ID, 'IMAGE', INSTRUCTION_GCS_URI)
 
     yield instruction
 
     # tear down
-    client = datalabeling.DataLabelingServiceClient()
+    testing_lib.delete_instruction(instruction.name)
 
-    # 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)
 
-    client.delete_instruction(instruction.name)
+@pytest.fixture(scope='module')
+def cleaner():
+    resource_names = []
+
+    yield resource_names
+
+    for resource_name in resource_names:
+        testing_lib.cancel_operation(resource_name)
 
 
 # Passing in dataset as the last argument in test_label_image since it needs
 # to be deleted before the annotation_spec_set can be deleted.
-@pytest.mark.flaky(max_runs=3)
-def test_label_text(capsys, annotation_spec_set, instruction, dataset):
-
-    # Start labeling.
-    response = label_text.label_text(
-        dataset.name,
-        instruction.name,
-        annotation_spec_set.name
-    )
+def test_label_text(capsys, annotation_spec_set, instruction, dataset, cleaner):
+
+    @backoff.on_exception(
+        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+    def run_sample():
+        # Start labeling.
+        return label_text.label_text(
+            dataset.name, instruction.name, annotation_spec_set.name)
+
+    response = run_sample()
+    cleaner.append(response.operation.name)
+
     out, _ = capsys.readouterr()
     assert 'Label_text operation name: ' in out
-    operation_name = response.operation.name
 
     # Cancels the labeling operation.
     response.cancel()
     assert response.cancelled() is True
-
-    client = datalabeling.DataLabelingServiceClient()
-
-    # 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)
-
-    client.transport._operations_client.cancel_operation(
-            operation_name)