Fix integration tests

Author: Balaji Veeramani

tests/integ/test_byo_estimator.py

@@ -20,6 +20,7 @@
 import sagemaker
 from sagemaker import image_uris
 from sagemaker.estimator import Estimator
+from sagemaker.serializers import BaseSerializer
 from sagemaker.utils import unique_name_from_base
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES, datasets
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
@@ -35,11 +36,15 @@ def training_set():
     return datasets.one_p_mnist()
 
 
-def fm_serializer(data):
-    js = {"instances": []}
-    for row in data:
-        js["instances"].append({"features": row.tolist()})
-    return json.dumps(js)
+class _FactorizationMachineSerializer(BaseSerializer):
+
+    CONTENT_TYPE = "application/json"
+
+    def serialize(data):
+        js = {"instances": []}
+        for row in data:
+            js["instances"].append({"features": row.tolist()})
+        return json.dumps(js)
 
 
 @pytest.mark.canary_quick
@@ -84,8 +89,7 @@ def test_byo_estimator(sagemaker_session, region, cpu_instance_type, training_se
     with timeout_and_delete_endpoint_by_name(job_name, sagemaker_session):
         model = estimator.create_model()
         predictor = model.deploy(1, cpu_instance_type, endpoint_name=job_name)
-        predictor.serializer = fm_serializer
-        predictor.content_type = "application/json"
+        predictor.serializer = _FactorizationMachineSerializer()
         predictor.deserializer = sagemaker.deserializers.JSONDeserializer()
 
         result = predictor.predict(training_set[0][:10])
@@ -130,8 +134,7 @@ def test_async_byo_estimator(sagemaker_session, region, cpu_instance_type, train
         )
         model = estimator.create_model()
         predictor = model.deploy(1, cpu_instance_type, endpoint_name=endpoint_name)
-        predictor.serializer = fm_serializer
-        predictor.content_type = "application/json"
+        predictor.serializer = _FactorizationMachineSerializer()
         predictor.deserializer = sagemaker.deserializers.JSONDeserializer()
 
         result = predictor.predict(training_set[0][:10])

tests/integ/test_tuner.py

@@ -28,6 +28,7 @@
 from sagemaker.estimator import Estimator
 from sagemaker.mxnet.estimator import MXNet
 from sagemaker.pytorch import PyTorch
+from sagemaker.serializers import BaseSerializer
 from sagemaker.tensorflow import TensorFlow
 from sagemaker.tuner import (
     IntegerParameter,
@@ -900,8 +901,7 @@ def test_tuning_byo_estimator(sagemaker_session, cpu_instance_type):
     best_training_job = tuner.best_training_job()
     with timeout_and_delete_endpoint_by_name(best_training_job, sagemaker_session):
         predictor = tuner.deploy(1, cpu_instance_type, endpoint_name=best_training_job)
-        predictor.serializer = _fm_serializer
-        predictor.content_type = "application/json"
+        predictor.serializer = _FactorizationMachineSerializer()
         predictor.deserializer = JSONDeserializer()
 
         result = predictor.predict(datasets.one_p_mnist()[0][:10])
@@ -912,11 +912,15 @@ def test_tuning_byo_estimator(sagemaker_session, cpu_instance_type):
 
 
 # Serializer for the Factorization Machines predictor (for BYO example)
-def _fm_serializer(data):
-    js = {"instances": []}
-    for row in data:
-        js["instances"].append({"features": row.tolist()})
-    return json.dumps(js)
+class _FactorizationMachineSerializer(BaseSerializer):
+
+    CONTENT_TYPE = "application/json"
+
+    def serializer(data):
+        js = {"instances": []}
+        for row in data:
+            js["instances"].append({"features": row.tolist()})
+        return json.dumps(js)
 
 
 def _assert_model_name_match(sagemaker_client, endpoint_config_name, model_name):

tests/integ/test_tuner_multi_algo.py

@@ -19,9 +19,9 @@
 
 from sagemaker import image_uris, utils
 from sagemaker.analytics import HyperparameterTuningJobAnalytics
-from sagemaker.content_types import CONTENT_TYPE_JSON
 from sagemaker.deserializers import JSONDeserializer
 from sagemaker.estimator import Estimator
+from sagemaker.serializers import BaseSerializer
 from sagemaker.tuner import ContinuousParameter, IntegerParameter, HyperparameterTuner
 from tests.integ import datasets, DATA_DIR, TUNING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
@@ -214,14 +214,17 @@ def _create_training_inputs(sagemaker_session):
 
 
 def _make_prediction(predictor, data):
-    predictor.serializer = _prediction_data_serializer
-    predictor.content_type = CONTENT_TYPE_JSON
+    predictor.serializer = PredictionDataSerializer()
     predictor.deserializer = JSONDeserializer()
     return predictor.predict(data)
 
 
-def _prediction_data_serializer(data):
-    js = {"instances": []}
-    for row in data:
-        js["instances"].append({"features": row.tolist()})
-    return json.dumps(js)
+class PredictionDataSerializer(BaseSerializer):
+
+    CONTENT_TYPE = "application/json"
+
+    def serialize(data):
+        js = {"instances": []}
+        for row in data:
+            js["instances"].append({"features": row.tolist()})
+        return json.dumps(js)