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

Takashi Matsuo · web-flow · commit 07bc5fc7e119 · 2020-05-15T11:19:19.000-07:00
fixes GoogleCloudPlatform#3760
diff --git a/datalabeling/create_annotation_spec_set_test.py b/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)
 
diff --git a/datalabeling/create_instruction_test.py b/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)
diff --git a/datalabeling/import_data_test.py b/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)
 
diff --git a/datalabeling/label_image_test.py b/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(
diff --git a/datalabeling/label_text_test.py b/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(
diff --git a/datalabeling/label_video_test.py b/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(
diff --git a/datalabeling/manage_dataset_test.py b/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)
 
