Fix unit test.

Author: Chuyang Deng

src/sagemaker/predictor.py

@@ -137,8 +137,8 @@ def delete_model(self):
     def _get_model_names(self):
         endpoint_desc = self.sagemaker_session.sagemaker_client.describe_endpoint(EndpointName=self.endpoint)
         endpoint_config_name = endpoint_desc['EndpointConfigName']
-        endpoint_config = self.sagemaker_session.sagemaker_client.describe_endpoint_config(EndpointConfigName=
-                                                                                               endpoint_config_name)
+        endpoint_config = self.sagemaker_session.sagemaker_client.describe_endpoint_config(
+            EndpointConfigName=endpoint_config_name)
         production_variants = endpoint_config['ProductionVariants']
         return map(lambda d: d['ModelName'], production_variants)
 

tests/unit/test_predictor.py

@@ -470,14 +470,17 @@ def test_delete_endpoint_only():
 
 def test_delete_model():
     endpoint_desc = {
+        'EndpointConfigName': 'my-endpoint-config'
+    }
+    endpoint_config_desc = {
         'ProductionVariants': [{
-            'VariantName': 'my-model'
+            'ModelName': 'my-model'
         }]
     }
     sagemaker_session = empty_sagemaker_session()
     sagemaker_session.sagemaker_client.describe_endpoint = Mock(return_value=endpoint_desc)
+    sagemaker_session.sagemaker_client.describe_endpoint_config = Mock(return_value=endpoint_config_desc)
     predictor = RealTimePredictor(ENDPOINT, sagemaker_session=sagemaker_session)
 
     predictor.delete_model()
     sagemaker_session.delete_model.assert_called_with('my-model')
-