Fix flaky for ml engine tests

Author: Jon Wayne Parrott

ml_engine/online_prediction/predict_test.py

@@ -15,6 +15,7 @@
 import json
 import socket
 
+from gcp_devrel.testing.flaky import flaky
 import pytest
 
 import predict
@@ -41,29 +42,29 @@
     BYTESTRING = f.read()
 
 
-@pytest.mark.flaky
+@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
+@flaky
 def test_predict_json_error():
     with pytest.raises(RuntimeError):
         predict.predict_json(
             PROJECT, MODEL, [{"foo": "bar"}], version=JSON_VERSION)
 
 
-@pytest.mark.flaky
+@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
+@flaky(max_runs=6)
 def test_predict_examples():
     result = predict.predict_examples(
         PROJECT, MODEL, [BYTESTRING, BYTESTRING], version=EXAMPLES_VERSION)

nox.py

@@ -136,7 +136,8 @@ def _setup_appengine_sdk(session):
 PYTEST_COMMON_ARGS = [
     '--cov',
     '--cov-config', os.path.abspath('.coveragerc'),
-    '--cov-report', 'term']
+    '--cov-report', 'term',
+    '--no-success-flaky-report']
 
 FLAKE8_COMMON_ARGS = [
     '--show-source', '--builtin', 'gettext', '--max-complexity', '20',
@@ -171,7 +172,7 @@ def _setup_appengine_sdk(session):
         NON_GAE_STANDARD_SAMPLES, CHANGED_FILES)
 
 
-def _session_tests(session, sample, post_install):
+def _session_tests(session, sample, post_install=None):
     """Runs py.test for a particular sample."""
     session.install('-r', 'testing/requirements.txt')
 

testing/requirements.txt

@@ -14,4 +14,4 @@ WebTest==2.0.25
 webapp2==2.5.2
 google-api-python-client==1.6.2
 google-cloud-core==0.23.0
-gcp-devrel-py-tools==0.0.1
+gcp-devrel-py-tools==0.0.4