Make predict test more lenient (GoogleCloudPlatform#4916)

Alec Glassford · vinnysenthil · web-flow · commit cd96bce47b01 · 2020-10-28T14:43:24.000-07:00
Co-authored-by: Vinny Senthil &lt;vinnysenthil@gmail.com&gt;
diff --git a/ml_engine/online_prediction/predict_test.py b/ml_engine/online_prediction/predict_test.py
@@ -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
