infra: add Model unit tests for delete_model and enable_network_isolation (#1424)

Author: laurenyu

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

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()