[datalabeling] testing: retry upon ServerError (#3762)

fixes #3760

Author: Takashi Matsuo

datalabeling/create_annotation_spec_set_test.py

@@ -17,7 +17,7 @@
 import os
 
 import backoff
-from google.api_core.exceptions import DeadlineExceeded
+from google.api_core.exceptions import ServerError
 import pytest
 
 import create_annotation_spec_set
@@ -40,7 +40,7 @@ def cleaner():
 def test_create_annotation_spec_set(cleaner, capsys):
 
     @backoff.on_exception(
-        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+        backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE)
     def run_sample():
         return create_annotation_spec_set.create_annotation_spec_set(PROJECT_ID)
 

datalabeling/create_instruction_test.py

@@ -17,7 +17,7 @@
 import os
 
 import backoff
-from google.api_core.exceptions import DeadlineExceeded
+from google.api_core.exceptions import ServerError
 import pytest
 
 import create_instruction
@@ -42,7 +42,7 @@ def cleaner():
 def test_create_instruction(cleaner, capsys):
 
     @backoff.on_exception(
-        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+        backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE)
     def run_sample():
         return create_instruction.create_instruction(
             PROJECT_ID, 'IMAGE', INSTRUCTION_GCS_URI)

datalabeling/import_data_test.py

@@ -17,7 +17,7 @@
 import os
 
 import backoff
-from google.api_core.exceptions import DeadlineExceeded
+from google.api_core.exceptions import ServerError
 import pytest
 
 import import_data
@@ -42,7 +42,7 @@ def dataset():
 def test_import_data(capsys, dataset):
 
     @backoff.on_exception(
-        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+        backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE)
     def run_sample():
         import_data.import_data(dataset.name, 'IMAGE', INPUT_GCS_URI)
 

datalabeling/label_image_test.py

@@ -17,7 +17,7 @@
 import os
 
 import backoff
-from google.api_core.exceptions import DeadlineExceeded
+from google.api_core.exceptions import ServerError
 import pytest
 
 import label_image
@@ -80,7 +80,7 @@ def test_label_image(
         capsys, annotation_spec_set, instruction, dataset, cleaner):
 
     @backoff.on_exception(
-        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+        backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE)
     def run_sample():
         # Start labeling.
         return label_image.label_image(

datalabeling/label_text_test.py

@@ -17,7 +17,7 @@
 import os
 
 import backoff
-from google.api_core.exceptions import DeadlineExceeded
+from google.api_core.exceptions import ServerError
 import pytest
 
 import label_text
@@ -80,7 +80,7 @@ def cleaner():
 def test_label_text(capsys, annotation_spec_set, instruction, dataset, cleaner):
 
     @backoff.on_exception(
-        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+        backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE)
     def run_sample():
         # Start labeling.
         return label_text.label_text(

datalabeling/label_video_test.py

@@ -17,7 +17,7 @@
 import os
 
 import backoff
-from google.api_core.exceptions import DeadlineExceeded
+from google.api_core.exceptions import ServerError
 import pytest
 
 import label_video
@@ -80,7 +80,7 @@ def test_label_video(
         capsys, annotation_spec_set, instruction, dataset, cleaner):
 
     @backoff.on_exception(
-        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+        backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE)
     def run_sample():
         # Start labeling.
         return label_video.label_video(

datalabeling/manage_dataset_test.py

@@ -17,8 +17,8 @@
 import os
 
 import backoff
-from google.api_core.exceptions import DeadlineExceeded
 from google.api_core.exceptions import RetryError
+from google.api_core.exceptions import ServerError
 import pytest
 
 import manage_dataset
@@ -60,7 +60,7 @@ def cleaner():
 def test_create_dataset(cleaner, capsys):
 
     @backoff.on_exception(
-        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+        backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE)
     def run_sample():
         return manage_dataset.create_dataset(PROJECT_ID)
 
@@ -75,7 +75,7 @@ def run_sample():
 def test_list_dataset(capsys, dataset):
 
     @backoff.on_exception(
-        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+        backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE)
     def run_sample():
         manage_dataset.list_datasets(PROJECT_ID)
 
@@ -86,7 +86,7 @@ def run_sample():
 
 def test_get_dataset(capsys, dataset):
     @backoff.on_exception(
-        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+        backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE)
     def run_sample():
         manage_dataset.get_dataset(dataset.name)
 
@@ -97,7 +97,7 @@ def run_sample():
 
 def test_delete_dataset(capsys, dataset):
     @backoff.on_exception(
-        backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE)
+        backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE)
     def run_sample():
         manage_dataset.delete_dataset(dataset.name)