add unit test for network isolation

Author: ajaykarpur

tests/unit/test_pipeline_model.py

@@ -298,3 +298,43 @@ def test_delete_model(tfo, time, sagemaker_session):
 
     pipeline_model.delete_model()
     sagemaker_session.delete_model.assert_called_with(pipeline_model.name)
+
+
+@patch("tarfile.open")
+@patch("time.strftime", return_value=TIMESTAMP)
+def test_network_isolation(tfo, time, sagemaker_session):
+    framework_model = DummyFrameworkModel(sagemaker_session)
+    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,
+        enable_network_isolation=True,
+    )
+    model.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=1)
+
+    sagemaker_session.create_model.assert_called_with(
+        model.name,
+        ROLE,
+        [
+            {
+                "Image": "mi-1",
+                "Environment": {
+                    "SAGEMAKER_PROGRAM": "blah.py",
+                    "SAGEMAKER_SUBMIT_DIRECTORY": "s3://mybucket/mi-1-2017-10-10-14-14-15/sourcedir.tar.gz",
+                    "SAGEMAKER_CONTAINER_LOG_LEVEL": "20",
+                    "SAGEMAKER_REGION": "us-west-2",
+                },
+                "ModelDataUrl": "s3://bucket/model_1.tar.gz",
+            },
+            {
+                "Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-sparkml-serving:2.2",
+                "Environment": {},
+                "ModelDataUrl": "s3://bucket/model_2.tar.gz",
+            },
+        ],
+        vpc_config=None,
+        enable_network_isolation=True,
+    )