Skip to content

Commit f877a9a

Browse files
author
Takashi Matsuo
authored
[datalabeling] testing: wrap rpcs with backoff (#3443)
* wrap all the rpcs with backoff * add a shared testing lib * remove flaky
1 parent c71f978 commit f877a9a

10 files changed

+316
-241
lines changed

datalabeling/create_annotation_spec_set_test.py

Lines changed: 27 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -16,29 +16,38 @@
1616

1717
import os
1818

19-
import create_annotation_spec_set
20-
from google.cloud import datalabeling_v1beta1 as datalabeling
21-
from google.api_core.client_options import ClientOptions
19+
import backoff
20+
from google.api_core.exceptions import DeadlineExceeded
2221
import pytest
2322

23+
import create_annotation_spec_set
24+
import testing_lib
25+
26+
2427
PROJECT_ID = os.getenv('GCLOUD_PROJECT')
2528

2629

27-
@pytest.mark.slow
28-
def test_create_annotation_spec_set(capsys):
29-
response = create_annotation_spec_set.create_annotation_spec_set(
30-
PROJECT_ID)
31-
out, _ = capsys.readouterr()
32-
assert 'The annotation_spec_set resource name:' in out
30+
@pytest.fixture(scope='module')
31+
def cleaner():
32+
resource_names = []
33+
34+
yield resource_names
3335

34-
# Delete the created annotation spec set.
35-
annotation_spec_set_name = response.name
36-
client = datalabeling.DataLabelingServiceClient()
36+
for resource_name in resource_names:
37+
testing_lib.delete_annotation_spec_set(resource_name)
3738

38-
# If provided, use a provided test endpoint - this will prevent tests on
39-
# this snippet from triggering any action by a real human
40-
if 'DATALABELING_ENDPOINT' in os.environ:
41-
opts = ClientOptions(api_endpoint=os.getenv('DATALABELING_ENDPOINT'))
42-
client = datalabeling.DataLabelingServiceClient(client_options=opts)
4339

44-
client.delete_annotation_spec_set(annotation_spec_set_name)
40+
def test_create_annotation_spec_set(cleaner, capsys):
41+
42+
@backoff.on_exception(
43+
backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
44+
def run_sample():
45+
return create_annotation_spec_set.create_annotation_spec_set(PROJECT_ID)
46+
47+
response = run_sample()
48+
49+
# For cleanup
50+
cleaner.append(response.name)
51+
52+
out, _ = capsys.readouterr()
53+
assert 'The annotation_spec_set resource name:' in out

datalabeling/create_instruction.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
import argparse
1818
import os
19+
1920
from google.api_core.client_options import ClientOptions
2021

2122

datalabeling/create_instruction_test.py

Lines changed: 26 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -16,34 +16,39 @@
1616

1717
import os
1818

19-
import create_instruction
20-
from google.api_core.client_options import ClientOptions
21-
from google.cloud import datalabeling_v1beta1 as datalabeling
19+
import backoff
20+
from google.api_core.exceptions import DeadlineExceeded
2221
import pytest
2322

23+
import create_instruction
24+
import testing_lib
25+
26+
2427
PROJECT_ID = os.getenv('GCLOUD_PROJECT')
2528
INSTRUCTION_GCS_URI = ('gs://cloud-samples-data/datalabeling'
2629
'/instruction/test.pdf')
2730

2831

29-
@pytest.mark.slow
30-
def test_create_instruction(capsys):
31-
result = create_instruction.create_instruction(
32-
PROJECT_ID,
33-
'IMAGE',
34-
INSTRUCTION_GCS_URI
35-
)
36-
out, _ = capsys.readouterr()
37-
assert 'The instruction resource name: ' in out
32+
@pytest.fixture(scope='module')
33+
def cleaner():
34+
resource_names = []
3835

39-
# Delete the created instruction.
40-
instruction_name = result.name
41-
client = datalabeling.DataLabelingServiceClient()
36+
yield resource_names
4237

43-
# If provided, use a provided test endpoint - this will prevent tests on
44-
# this snippet from triggering any action by a real human
45-
if 'DATALABELING_ENDPOINT' in os.environ:
46-
opts = ClientOptions(api_endpoint=os.getenv('DATALABELING_ENDPOINT'))
47-
client = datalabeling.DataLabelingServiceClient(client_options=opts)
38+
for resource_name in resource_names:
39+
testing_lib.delete_instruction(resource_name)
4840

49-
client.delete_instruction(instruction_name)
41+
42+
def test_create_instruction(cleaner, capsys):
43+
44+
@backoff.on_exception(
45+
backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
46+
def run_sample():
47+
return create_instruction.create_instruction(
48+
PROJECT_ID, 'IMAGE', INSTRUCTION_GCS_URI)
49+
50+
instruction = run_sample()
51+
cleaner.append(instruction.name)
52+
53+
out, _ = capsys.readouterr()
54+
assert 'The instruction resource name: ' in out

datalabeling/import_data_test.py

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,27 +16,36 @@
1616

1717
import os
1818

19-
import import_data
20-
import manage_dataset
19+
import backoff
20+
from google.api_core.exceptions import DeadlineExceeded
2121
import pytest
2222

23+
import import_data
24+
import testing_lib
25+
26+
2327
PROJECT_ID = os.getenv('GCLOUD_PROJECT')
2428
INPUT_GCS_URI = 'gs://cloud-samples-data/datalabeling/image/image_dataset.csv'
2529

2630

27-
@pytest.fixture(scope='function')
31+
@pytest.fixture(scope='module')
2832
def dataset():
2933
# create a temporary dataset
30-
dataset = manage_dataset.create_dataset(PROJECT_ID)
34+
dataset = testing_lib.create_dataset(PROJECT_ID)
3135

3236
yield dataset
3337

3438
# tear down
35-
manage_dataset.delete_dataset(dataset.name)
39+
testing_lib.delete_dataset(dataset.name)
3640

3741

38-
@pytest.mark.slow
3942
def test_import_data(capsys, dataset):
40-
import_data.import_data(dataset.name, 'IMAGE', INPUT_GCS_URI)
43+
44+
@backoff.on_exception(
45+
backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
46+
def run_sample():
47+
import_data.import_data(dataset.name, 'IMAGE', INPUT_GCS_URI)
48+
49+
run_sample()
4150
out, _ = capsys.readouterr()
4251
assert 'Dataset resource name: ' in out

datalabeling/label_image_test.py

Lines changed: 39 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -16,100 +16,78 @@
1616

1717
import os
1818

19-
import create_annotation_spec_set
20-
import create_instruction
21-
from google.api_core.client_options import ClientOptions
22-
from google.cloud import datalabeling_v1beta1 as datalabeling
23-
import import_data
24-
import label_image
25-
import manage_dataset
19+
import backoff
20+
from google.api_core.exceptions import DeadlineExceeded
2621
import pytest
2722

23+
import label_image
24+
import testing_lib
25+
26+
2827
PROJECT_ID = os.getenv('GCLOUD_PROJECT')
2928
INPUT_GCS_URI = 'gs://cloud-samples-data/datalabeling/image/image_dataset.csv'
29+
INSTRUCTION_GCS_URI = ('gs://cloud-samples-data/datalabeling'
30+
'/instruction/test.pdf')
3031

3132

32-
@pytest.fixture(scope='function')
33+
@pytest.fixture(scope='module')
3334
def dataset():
3435
# create a temporary dataset
35-
dataset = manage_dataset.create_dataset(PROJECT_ID)
36-
37-
# import some data to it
38-
import_data.import_data(dataset.name, 'IMAGE', INPUT_GCS_URI)
36+
dataset = testing_lib.create_dataset(PROJECT_ID)
3937

38+
testing_lib.import_data(dataset.name, 'IMAGE', INPUT_GCS_URI)
4039
yield dataset
4140

4241
# tear down
43-
manage_dataset.delete_dataset(dataset.name)
42+
testing_lib.delete_dataset(dataset.name)
4443

4544

46-
@pytest.fixture(scope='function')
45+
@pytest.fixture(scope='module')
4746
def annotation_spec_set():
4847
# create a temporary annotation_spec_set
49-
response = create_annotation_spec_set.create_annotation_spec_set(
50-
PROJECT_ID)
48+
response = testing_lib.create_annotation_spec_set(PROJECT_ID)
5149

5250
yield response
5351

54-
# tear down
55-
client = datalabeling.DataLabelingServiceClient()
56-
57-
# If provided, use a provided test endpoint - this will prevent tests on
58-
# this snippet from triggering any action by a real human
59-
if 'DATALABELING_ENDPOINT' in os.environ:
60-
opts = ClientOptions(api_endpoint=os.getenv('DATALABELING_ENDPOINT'))
61-
client = datalabeling.DataLabelingServiceClient(client_options=opts)
62-
63-
client.delete_annotation_spec_set(response.name)
52+
testing_lib.delete_annotation_spec_set(response.name)
6453

6554

66-
@pytest.fixture(scope='function')
55+
@pytest.fixture(scope='module')
6756
def instruction():
6857
# create a temporary instruction
69-
instruction = create_instruction.create_instruction(
70-
PROJECT_ID, 'IMAGE',
71-
'gs://cloud-samples-data/datalabeling/instruction/test.pdf')
58+
instruction = testing_lib.create_instruction(
59+
PROJECT_ID, 'IMAGE', INSTRUCTION_GCS_URI)
7260

7361
yield instruction
7462

7563
# tear down
76-
client = datalabeling.DataLabelingServiceClient()
64+
testing_lib.delete_instruction(instruction.name)
7765

78-
# If provided, use a provided test endpoint - this will prevent tests on
79-
# this snippet from triggering any action by a real human
80-
if 'DATALABELING_ENDPOINT' in os.environ:
81-
opts = ClientOptions(api_endpoint=os.getenv('DATALABELING_ENDPOINT'))
82-
client = datalabeling.DataLabelingServiceClient(client_options=opts)
8366

84-
client.delete_instruction(instruction.name)
67+
@pytest.fixture(scope='module')
68+
def cleaner():
69+
resource_names = []
70+
71+
yield resource_names
72+
73+
for resource_name in resource_names:
74+
testing_lib.cancel_operation(resource_name)
8575

8676

8777
# Passing in dataset as the last argument in test_label_image since it needs
8878
# to be deleted before the annotation_spec_set can be deleted.
89-
@pytest.mark.slow
90-
def test_label_image(capsys, annotation_spec_set, instruction, dataset):
91-
92-
# Start labeling.
93-
response = label_image.label_image(
94-
dataset.name,
95-
instruction.name,
96-
annotation_spec_set.name
97-
)
98-
out, _ = capsys.readouterr()
99-
assert 'Label_image operation name: ' in out
100-
operation_name = response.operation.name
79+
def test_label_image(
80+
capsys, annotation_spec_set, instruction, dataset, cleaner):
10181

102-
# Cancels the labeling operation.
103-
response.cancel()
104-
assert response.cancelled() is True
82+
@backoff.on_exception(
83+
backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
84+
def run_sample():
85+
# Start labeling.
86+
return label_image.label_image(
87+
dataset.name, instruction.name, annotation_spec_set.name)
10588

106-
client = datalabeling.DataLabelingServiceClient()
89+
response = run_sample()
90+
cleaner.append(response.operation.name)
10791

108-
# If provided, use a provided test endpoint - this will prevent tests on
109-
# this snippet from triggering any action by a real human
110-
if 'DATALABELING_ENDPOINT' in os.environ:
111-
opts = ClientOptions(api_endpoint=os.getenv('DATALABELING_ENDPOINT'))
112-
client = datalabeling.DataLabelingServiceClient(client_options=opts)
113-
114-
client.transport._operations_client.cancel_operation(
115-
operation_name)
92+
out, _ = capsys.readouterr()
93+
assert 'Label_image operation name: ' in out

0 commit comments

Comments
 (0)