code4ward
diff --git a/‎datalabeling/create_annotation_spec_set_test.py‎
Lines changed: 27 additions & 18 deletions b/‎datalabeling/create_annotation_spec_set_test.py‎
Lines changed: 27 additions & 18 deletions
diff --git a/‎datalabeling/create_instruction.py‎
Lines changed: 1 addition & 0 deletions b/‎datalabeling/create_instruction.py‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎datalabeling/create_instruction_test.py‎
Lines changed: 26 additions & 21 deletions b/‎datalabeling/create_instruction_test.py‎
Lines changed: 26 additions & 21 deletions
diff --git a/‎datalabeling/import_data_test.py‎
Lines changed: 16 additions & 7 deletions b/‎datalabeling/import_data_test.py‎
Lines changed: 16 additions & 7 deletions
diff --git a/‎datalabeling/label_image_test.py‎
Lines changed: 39 additions & 61 deletions b/‎datalabeling/label_image_test.py‎
Lines changed: 39 additions & 61 deletions
@@ -16,29 +16,38 @@
 
 import os
 
-import create_annotation_spec_set
-from google.cloud import datalabeling_v1beta1 as datalabeling
-from google.api_core.client_options import ClientOptions
+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.slow
-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
+@pytest.fixture(scope='module')
+def cleaner():
+    resource_names = []
+
+    yield resource_names
 
-    # Delete the created annotation spec set.
-    annotation_spec_set_name = response.name
-    client = datalabeling.DataLabelingServiceClient()
+    for resource_name in resource_names:
+        testing_lib.delete_annotation_spec_set(resource_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_annotation_spec_set(annotation_spec_set_name)
+def test_create_annotation_spec_set(cleaner, capsys):
+
+    @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)
+
+    response = run_sample()
+
+    # For cleanup
+    cleaner.append(response.name)
+
+    out, _ = capsys.readouterr()
+    assert 'The annotation_spec_set resource name:' in out
@@ -16,6 +16,7 @@
 
 import argparse
 import os
+
 from google.api_core.client_options import ClientOptions
 
 
 
@@ -16,34 +16,39 @@
 
 import os
 
-import create_instruction
-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.slow
-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 = []
 
-    # Delete the created instruction.
-    instruction_name = result.name
-    client = datalabeling.DataLabelingServiceClient()
+    yield resource_names
 
-    # 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)
+    for resource_name in resource_names:
+        testing_lib.delete_instruction(resource_name)
 
-    client.delete_instruction(instruction_name)
+
+def test_create_instruction(cleaner, capsys):
+
+    @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
@@ -16,27 +16,36 @@
 
 import os
 
-import import_data
-import manage_dataset
+import backoff
+from google.api_core.exceptions import DeadlineExceeded
 import pytest
 
+import import_data
+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')
+@pytest.fixture(scope='module')
 def dataset():
     # create a temporary dataset
-    dataset = manage_dataset.create_dataset(PROJECT_ID)
+    dataset = testing_lib.create_dataset(PROJECT_ID)
 
     yield dataset
 
     # tear down
-    manage_dataset.delete_dataset(dataset.name)
+    testing_lib.delete_dataset(dataset.name)
 
 
-@pytest.mark.slow
 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
@@ -16,100 +16,78 @@
 
 import os
 
-import create_annotation_spec_set
-import create_instruction
-from google.api_core.client_options import ClientOptions
-from google.cloud import datalabeling_v1beta1 as datalabeling
-import import_data
-import label_image
-import manage_dataset
+import backoff
+from google.api_core.exceptions import DeadlineExceeded
 import pytest
 
+import label_image
+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')
+@pytest.fixture(scope='module')
 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)
 
+    testing_lib.import_data(dataset.name, 'IMAGE', INPUT_GCS_URI)
     yield dataset
 
     # tear down
-    manage_dataset.delete_dataset(dataset.name)
+    testing_lib.delete_dataset(dataset.name)
 
 
-@pytest.fixture(scope='function')
+@pytest.fixture(scope='module')
 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')
+@pytest.fixture(scope='module')
 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)
 
-    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.slow
-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