Fix py2-3 compat issue in sending bytestring (GoogleCloudPlatform#878)

elibixby · dpebot · commit 559c791dd711 · 2017-03-24T14:58:02.000-07:00
* Fix py2-3 compat issue in sending bytestring

* unmark test_predict_examples as slow

* Move data to resources

* Make test_census_example_bytes non-trivial

* Fix linter

* b64 encoded string must be decoded to utf-8
diff --git a/ml_engine/online_prediction/predict.py b/ml_engine/online_prediction/predict.py
@@ -89,7 +89,7 @@ def predict_examples(project,
     response = service.projects().predict(
         name=name,
         body={'instances': [
-            {'b64': base64.b64encode(example_bytes)}
+            {'b64': base64.b64encode(example_bytes).decode('utf-8')}
             for example_bytes in example_bytes_list
         ]}
     ).execute()
diff --git a/ml_engine/online_prediction/predict_test.py b/ml_engine/online_prediction/predict_test.py
@@ -14,7 +14,7 @@
 
 """Tests for predict.py ."""
 
-import base64
+import json
 
 import pytest
 
@@ -25,27 +25,20 @@
 JSON_VERSION = 'v1json'
 EXAMPLES_VERSION = 'v1example'
 PROJECT = 'python-docs-samples-tests'
-JSON = {
-    'age': 25,
-    'workclass': ' Private',
-    'education': ' 11th',
-    'education_num': 7,
-    'marital_status': ' Never-married',
-    'occupation': ' Machine-op-inspct',
-    'relationship': ' Own-child',
-    'race': ' Black',
-    'gender': ' Male',
-    'capital_gain': 0,
-    'capital_loss': 0,
-    'hours_per_week': 40,
-    'native_country': ' United-States'
-}
 EXPECTED_OUTPUT = {
     u'confidence': 0.7760371565818787,
     u'predictions': u' <=50K'
 }
 
 
+with open('resources/census_test_data.json') as f:
+    JSON = json.load(f)
+
+
+with open('resources/census_example_bytes.pb', 'rb') as f:
+    BYTESTRING = f.read()
+
+
 def test_predict_json():
     result = predict.predict_json(
         PROJECT, MODEL, [JSON, JSON], version=JSON_VERSION)
@@ -60,13 +53,13 @@ def test_predict_json_error():
 
 @pytest.mark.slow
 def test_census_example_to_bytes():
+    import tensorflow as tf
     b = predict.census_to_example_bytes(JSON)
-    assert base64.b64encode(b) is not None
+    assert tf.train.Example.FromString(b) == tf.train.Example.FromString(
+        BYTESTRING)
 
 
-@pytest.mark.slow
 def test_predict_examples():
-    b = predict.census_to_example_bytes(JSON)
     result = predict.predict_examples(
-        PROJECT, MODEL, [b, b], version=EXAMPLES_VERSION)
+        PROJECT, MODEL, [BYTESTRING, BYTESTRING], version=EXAMPLES_VERSION)
     assert [EXPECTED_OUTPUT, EXPECTED_OUTPUT] == result
diff --git a/ml_engine/online_prediction/resources/census_example_bytes.pb b/ml_engine/online_prediction/resources/census_example_bytes.pb
diff --git a/ml_engine/online_prediction/resources/census_test_data.json b/ml_engine/online_prediction/resources/census_test_data.json
@@ -0,0 +1 @@
+{"hours_per_week": 40, "native_country": " United-States", "relationship": " Own-child", "capital_loss": 0, "education": " 11th", "capital_gain": 0, "occupation": " Machine-op-inspct", "workclass": " Private", "gender": " Male", "age": 25, "marital_status": " Never-married", "race": " Black", "education_num": 7}

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+{"hours_per_week": 40, "native_country": " United-States", "relationship": " Own-child", "capital_loss": 0, "education": " 11th", "capital_gain": 0, "occupation": " Machine-op-inspct", "workclass": " Private", "gender": " Male", "age": 25, "marital_status": " Never-married", "race": " Black", "education_num": 7}`