add unit and integration tests

Author: imujjwal96

tests/integ/test_inference_pipeline.py

@@ -144,3 +144,53 @@ def test_inference_pipeline_model_deploy(sagemaker_session):
     with pytest.raises(Exception) as exception:
         sagemaker_session.sagemaker_client.describe_model(ModelName=model.name)
         assert "Could not find model" in str(exception.value)
+
+
+def test_inference_pipeline_model_deploy_with_update_endpoint(sagemaker_session):
+    sparkml_data_path = os.path.join(DATA_DIR, "sparkml_model")
+    xgboost_data_path = os.path.join(DATA_DIR, "xgboost_model")
+    endpoint_name = "test-inference-pipeline-deploy-{}".format(sagemaker_timestamp())
+    sparkml_model_data = sagemaker_session.upload_data(
+        path=os.path.join(sparkml_data_path, "mleap_model.tar.gz"),
+        key_prefix="integ-test-data/sparkml/model",
+    )
+    xgb_model_data = sagemaker_session.upload_data(
+        path=os.path.join(xgboost_data_path, "xgb_model.tar.gz"),
+        key_prefix="integ-test-data/xgboost/model",
+    )
+
+    with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+        sparkml_model = SparkMLModel(
+            model_data=sparkml_model_data,
+            env={"SAGEMAKER_SPARKML_SCHEMA": SCHEMA},
+            sagemaker_session=sagemaker_session,
+        )
+        xgb_image = get_image_uri(sagemaker_session.boto_region_name, "xgboost")
+        xgb_model = Model(
+            model_data=xgb_model_data, image=xgb_image, sagemaker_session=sagemaker_session
+        )
+        model = PipelineModel(
+            models=[sparkml_model, xgb_model],
+            role="SageMakerRole",
+            sagemaker_session=sagemaker_session,
+            name=endpoint_name,
+        )
+        model.deploy(1, "ml.m4.xlarge", endpoint_name=endpoint_name)
+        old_endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)
+        old_config_name = old_endpoint["EndpointConfigName"]
+
+        model.deploy(1, "ml.m4.xlarge", update_endpoint=True, endpoint_name=endpoint_name)
+        new_endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)[
+            "ProductionVariants"
+        ]
+        new_production_variants = new_endpoint["ProductionVariants"]
+        new_config_name = new_endpoint["EndpointConfigName"]
+
+        assert old_config_name != new_config_name
+        assert new_production_variants["InstanceType"] == "ml.m4.xlarge"
+        assert new_production_variants["InitialInstanceCount"] == 1
+
+    model.delete_model()
+    with pytest.raises(Exception) as exception:
+        sagemaker_session.sagemaker_client.describe_model(ModelName=model.name)
+        assert "Could not find model" in str(exception.value)

tests/unit/test_pipeline_model.py

@@ -159,6 +159,41 @@ def test_deploy_endpoint_name(tfo, time, sagemaker_session):
     )
 
 
+@patch("tarfile.open")
+@patch("time.strftime", return_value=TIMESTAMP)
+def test_deploy_update_endpoint(tfo, time, sagemaker_session):
+    framework_model = DummyFrameworkModel(sagemaker_session)
+    endpoint_name = "endpoint-name"
+    sparkml_model = SparkMLModel(
+        model_data=MODEL_DATA_2, role=ROLE, sagemaker_session=sagemaker_session
+    )
+    model = PipelineModel(
+        models=[framework_model, sparkml_model], role=ROLE, sagemaker_session=sagemaker_session
+    )
+    model.deploy(
+        instance_type=INSTANCE_TYPE,
+        initial_instance_count=1,
+        endpoint_name=endpoint_name,
+        update_endpoint=True,
+    )
+
+    sagemaker_session.create_endpoint_config.assert_called_with(
+        name=model.name,
+        model_name=model.name,
+        initial_instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+        tags=None,
+    )
+    config_name = sagemaker_session.create_endpoint_config(
+        name=model.name,
+        model_name=model.name,
+        initial_instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+    )
+    sagemaker_session.update_endpoint.assert_called_with(endpoint_name, config_name)
+    sagemaker_session.create_endpoint.assert_not_called()
+
+
 @patch("tarfile.open")
 @patch("time.strftime", return_value=TIMESTAMP)
 def test_transformer(tfo, time, sagemaker_session):