Mark ml engine tests as flaky

Jon Wayne Parrott · Jon Wayne Parrott · commit 536baf654d02 · 2017-04-06T10:13:13.000-07:00
diff --git a/ml_engine/online_prediction/predict_test.py b/ml_engine/online_prediction/predict_test.py
@@ -12,15 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for predict.py ."""
-
 import json
+import socket
 
 import pytest
 
 import predict
 
-
 MODEL = 'census'
 JSON_VERSION = 'v1json'
 EXAMPLES_VERSION = 'v1example'
@@ -30,6 +28,10 @@
     u'predictions': u' <=50K'
 }
 
+# Raise the socket timeout. The requests involved in the sample can take
+# a long time to complete.
+socket.setdefaulttimeout(60)
+
 
 with open('resources/census_test_data.json') as f:
     JSON = json.load(f)
@@ -39,26 +41,29 @@
     BYTESTRING = f.read()
 
 
+@pytest.mark.flaky
 def test_predict_json():
     result = predict.predict_json(
         PROJECT, MODEL, [JSON, JSON], version=JSON_VERSION)
     assert [EXPECTED_OUTPUT, EXPECTED_OUTPUT] == result
 
 
+@pytest.mark.flaky
 def test_predict_json_error():
     with pytest.raises(RuntimeError):
         predict.predict_json(
             PROJECT, MODEL, [{"foo": "bar"}], version=JSON_VERSION)
 
 
-@pytest.mark.slow
+@pytest.mark.flaky
 def test_census_example_to_bytes():
     import tensorflow as tf
     b = predict.census_to_example_bytes(JSON)
     assert tf.train.Example.FromString(b) == tf.train.Example.FromString(
         BYTESTRING)
 
 
+@pytest.mark.flaky
 def test_predict_examples():
     result = predict.predict_examples(
         PROJECT, MODEL, [BYTESTRING, BYTESTRING], version=EXAMPLES_VERSION)
