test: add retries to automl language test (#8870)

GoogleCloudPlatform · Jan 5, 2023 · 32c74e3 · 32c74e3
1 parent ee28639
commit 32c74e3
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/automl/snippets/language_entity_extraction_predict_test.py b/automl/snippets/language_entity_extraction_predict_test.py
@@ -14,6 +14,9 @@
 
 import os
 
+import backoff
+from google.api_core.exceptions import InternalServerError
+from google.api_core.exceptions import ServiceUnavailable
 from google.cloud import automl
 import pytest
 
@@ -35,6 +38,7 @@ def verify_model_state():
         response.result()
 
 
+@backoff.on_exception(backoff.expo, (InternalServerError, ServiceUnavailable), max_tries=3)
 def test_predict(capsys, verify_model_state):
     verify_model_state
     text = (