Fix tests

Author: Balaji Veeramani

src/sagemaker/multidatamodel.py

@@ -211,12 +211,12 @@ def deploy(
             enable_network_isolation = self.model.enable_network_isolation()
             role = self.model.role
             vpc_config = self.model.vpc_config
-            predictor = self.model.predictor_cls
+            predictor_cls = self.model.predictor_cls
         else:
             enable_network_isolation = self.enable_network_isolation()
             role = self.role
             vpc_config = self.vpc_config
-            predictor = self.predictor_cls
+            predictor_cls = self.predictor_cls
 
         if role is None:
             raise ValueError("Role can not be null for deploying a model")
@@ -255,8 +255,8 @@ def deploy(
             data_capture_config_dict=data_capture_config_dict,
         )
 
-        if predictor:
-            predictor = self.predictor_cls(self.endpoint_name, self.sagemaker_session)
+        if predictor_cls:
+            predictor = predictor_cls(self.endpoint_name, self.sagemaker_session)
             if serializer:
                 predictor.serializer = serializer
             if deserializer:

src/sagemaker/tensorflow/model.py

@@ -196,6 +196,8 @@ def deploy(
         self,
         initial_instance_count,
         instance_type,
+        serializer=None,
+        deserializer=None,
         accelerator_type=None,
         endpoint_name=None,
         tags=None,
@@ -211,6 +213,8 @@ def deploy(
         return super(TensorFlowModel, self).deploy(
             initial_instance_count=initial_instance_count,
             instance_type=instance_type,
+            serializer=serializer,
+            deserializer=deserializer,
             accelerator_type=accelerator_type,
             endpoint_name=endpoint_name,
             tags=tags,

tests/unit/sagemaker/automl/test_auto_ml.py

@@ -591,6 +591,8 @@ def test_deploy_optional_args(candidate_estimator, sagemaker_session, candidate_
     mock_pipeline.deploy.assert_called_with(
         initial_instance_count=INSTANCE_COUNT,
         instance_type=INSTANCE_TYPE,
+        serializer=None,
+        deserializer=None,
         endpoint_name=JOB_NAME,
         tags=TAGS,
         wait=False,

tests/unit/test_estimator.py

@@ -2122,7 +2122,7 @@ def test_generic_deploy_accelerator_type(sagemaker_session):
         IMAGE_URI, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, sagemaker_session=sagemaker_session
     )
     e.fit({"train": "s3://bucket/training-prefix"})
-    e.deploy(INSTANCE_COUNT, INSTANCE_TYPE, ACCELERATOR_TYPE)
+    e.deploy(INSTANCE_COUNT, INSTANCE_TYPE, accelerator_type=ACCELERATOR_TYPE)
 
     args = e.sagemaker_session.endpoint_from_production_variants.call_args[1]
     print(args)

tests/unit/test_tuner.py

@@ -893,6 +893,8 @@ def test_deploy_optional_params(_get_best_training_job, best_estimator, tuner):
     estimator.deploy.assert_called_with(
         initial_instance_count=INSTANCE_COUNT,
         instance_type=INSTANCE_TYPE,
+        serializer=None,
+        deserializer=None,
         accelerator_type=accelerator,
         endpoint_name=endpoint_name,
         wait=False,