test: Backoff on predict_test.py (#302)

* docs: removed default model comment and added comment to test

* fix: backoff on predict_test.py

* fix: backoff on predict_test.py

* fix: added backoff to predict_test.py

* fix: added backoff to predict_test.py

* fix: added backoff to predict_test.py

Author: nayaknishant

automl/snippets/translate_create_model.py

@@ -27,7 +27,6 @@ def create_model(project_id, dataset_id, display_name):
 
     # A resource that represents Google Cloud Platform location.
     project_location = f"projects/{project_id}/locations/us-central1"
-    # Leave model unset to use the default base model provided by Google
     translation_model_metadata = automl.TranslationModelMetadata()
     model = automl.Model(
         display_name=display_name,

automl/snippets/translate_create_model_test.py

@@ -26,6 +26,10 @@ def test_translate_create_model(capsys):
             PROJECT_ID, DATASET_ID, "translate_test_create_model"
         )
         out, _ = capsys.readouterr()
+        # After setting DATASET_ID, change line below to
+        # assert "Training started..." in out
         assert "Dataset does not exist." in out
     except Exception as e:
+        # After setting DATASET_ID, change line below to
+        # assert "Training started..." in e.message
         assert "Dataset does not exist." in e.message

automl/tables/predict_test.py

@@ -16,18 +16,23 @@
 
 import os
 
+import backoff
+
 from google.cloud.automl_v1beta1 import Model
 
 import automl_tables_model
 import automl_tables_predict
 import model_test
 
-
 PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
 REGION = "us-central1"
 STATIC_MODEL = model_test.STATIC_MODEL
+MAX_TIMEOUT = 200
 
 
+@backoff.on_exception(
+    wait_gen=lambda: iter([50, 150, MAX_TIMEOUT]), exception=Exception, max_tries=3
+)
 def test_predict(capsys):
     inputs = {
         "Age": 31,

automl/tables/requirements-test.txt

@@ -1 +1,2 @@
 pytest==6.2.5
+backoff==1.11.1