From 58352ab45bd2136449cf8dbe1f11ce463afaa5f0 Mon Sep 17 00:00:00 2001
From: Lauren Yu <6631887+laurenyu@users.noreply.github.com>
Date: Thu, 16 Apr 2020 08:50:39 -0700
Subject: [PATCH] infra: add Model unit tests for delete_model and
 enable_network_isolation

---
 .../sagemaker/model/test_framework_model.py   | 23 -----------------
 tests/unit/sagemaker/model/test_model.py      | 25 +++++++++++++++++++
 2 files changed, 25 insertions(+), 23 deletions(-)

diff --git a/tests/unit/sagemaker/model/test_framework_model.py b/tests/unit/sagemaker/model/test_framework_model.py
index 120594d9c3..c3498a839e 100644
--- a/tests/unit/sagemaker/model/test_framework_model.py
+++ b/tests/unit/sagemaker/model/test_framework_model.py
@@ -237,29 +237,6 @@ def test_deploy_update_endpoint_optional_args(sagemaker_session, tmpdir):
     sagemaker_session.create_endpoint.assert_not_called()
 
 
-def test_model_enable_network_isolation(sagemaker_session):
-    model = DummyFrameworkModel(sagemaker_session=sagemaker_session)
-    assert model.enable_network_isolation() is False
-
-
-@patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
-@patch("time.strftime", MagicMock(return_value=TIMESTAMP))
-def test_model_delete_model(sagemaker_session, tmpdir):
-    model = DummyFrameworkModel(sagemaker_session, source_dir=str(tmpdir))
-    model.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=1)
-    model.delete_model()
-
-    sagemaker_session.delete_model.assert_called_with(model.name)
-
-
-def test_delete_non_deployed_model(sagemaker_session):
-    model = DummyFrameworkModel(sagemaker_session)
-    with pytest.raises(
-        ValueError, match="The SageMaker model must be created first before attempting to delete."
-    ):
-        model.delete_model()
-
-
 def test_compile_model_for_inferentia(sagemaker_session, tmpdir):
     sagemaker_session.wait_for_compilation_job = Mock(
         return_value=DESCRIBE_COMPILATION_JOB_RESPONSE
diff --git a/tests/unit/sagemaker/model/test_model.py b/tests/unit/sagemaker/model/test_model.py
index 8975849be3..4608de7994 100644
--- a/tests/unit/sagemaker/model/test_model.py
+++ b/tests/unit/sagemaker/model/test_model.py
@@ -53,6 +53,14 @@ def test_prepare_container_def():
     assert expected == container_def
 
 
+def test_model_enable_network_isolation():
+    model = Model(MODEL_DATA, MODEL_IMAGE)
+    assert model.enable_network_isolation() is False
+
+    model = Model(MODEL_DATA, MODEL_IMAGE, enable_network_isolation=True)
+    assert model.enable_network_isolation()
+
+
 @patch("sagemaker.model.Model.prepare_container_def")
 @patch("sagemaker.utils.name_from_image")
 def test_create_sagemaker_model(name_from_image, prepare_container_def, sagemaker_session):
@@ -468,3 +476,20 @@ def test_transformer_creates_correct_session(local_session, session):
     transformer = model.transformer(instance_count=1, instance_type="ml.m5.xlarge")
     assert model.sagemaker_session == session.return_value
     assert transformer.sagemaker_session == session.return_value
+
+
+def test_delete_model(sagemaker_session):
+    model = Model(MODEL_DATA, MODEL_IMAGE, name=MODEL_NAME, sagemaker_session=sagemaker_session)
+
+    model.delete_model()
+    sagemaker_session.delete_model.assert_called_with(model.name)
+
+
+def test_delete_model_no_name(sagemaker_session):
+    model = Model(MODEL_DATA, MODEL_IMAGE, sagemaker_session=sagemaker_session)
+
+    with pytest.raises(
+        ValueError, match="The SageMaker model must be created first before attempting to delete."
+    ):
+        model.delete_model()
+    sagemaker_session.delete_model.assert_not_called()