mywbob
diff --git a/‎ml_engine/online_prediction/predict_test.py
Copy file name to clipboardExpand all lines: ml_engine/online_prediction/predict_test.py
+13-3Lines changed: 13 additions & 3 deletions b/‎ml_engine/online_prediction/predict_test.py
Copy file name to clipboardExpand all lines: ml_engine/online_prediction/predict_test.py
+13-3Lines changed: 13 additions & 3 deletions
@@ -22,10 +22,13 @@
 MODEL = 'census'
 JSON_VERSION = 'v2json'
 PROJECT = 'python-docs-samples-tests'
+CONF_KEY = u'confidence'
+PRED_KEY = u'predictions'
 EXPECTED_OUTPUT = {
-    u'confidence': 0.7760370969772339,
-    u'predictions': u' <=50K'
+    CONF_KEY: 0.7760370969772339,
+    PRED_KEY: u' <=50K'
 }
+CONFIDENCE_EPSILON = 1e-4
 
 # Raise the socket timeout. The requests involved in the sample can take
 # a long time to complete.
@@ -40,7 +43,14 @@
 def test_predict_json():
     result = predict.predict_json(
         PROJECT, MODEL, [JSON, JSON], version=JSON_VERSION)
-    assert [EXPECTED_OUTPUT, EXPECTED_OUTPUT] == result
+    # Result contains two identical predictions
+    assert len(result) == 2 and result[0] == result[1]
+    # Each prediction has `confidence` and `predictions`
+    assert result[0].keys() == EXPECTED_OUTPUT.keys()
+    # Prediction matches
+    assert result[0][PRED_KEY] == EXPECTED_OUTPUT[PRED_KEY]
+    # Confidence within epsilon
+    assert abs(result[0][CONF_KEY] - EXPECTED_OUTPUT[CONF_KEY]) < CONFIDENCE_EPSILON
 
 
 @pytest.mark.flaky