change: move unit tests

Author: evakravi

tests/unit/sagemaker/model/test_model.py

@@ -11,12 +11,13 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
+from unittest.mock import MagicMock
 
 import pytest
 from mock import Mock, patch
 
 import sagemaker
-from sagemaker.model import Model
+from sagemaker.model import FrameworkModel, Model
 
 MODEL_DATA = "s3://bucket/model.tar.gz"
 MODEL_IMAGE = "mi"
@@ -27,10 +28,39 @@
 INSTANCE_TYPE = "ml.c4.4xlarge"
 ROLE = "some-role"
 
+REGION = "us-west-2"
+BUCKET_NAME = "some-bucket-name"
+GIT_REPO = "https://github.com/aws/sagemaker-python-sdk.git"
+BRANCH = "test-branch-git-config"
+COMMIT = "ae15c9d7d5b97ea95ea451e4662ee43da3401d73"
+ENTRY_POINT_INFERENCE = "inference.py"
 
-@pytest.fixture
+SCRIPT_URI = "s3://codebucket/someprefix/sourcedir.tar.gz"
+IMAGE_URI = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.9.0-gpu-py38"
+
+
+class DummyFrameworkModel(FrameworkModel):
+    def __init__(self, **kwargs):
+        super(DummyFrameworkModel, self).__init__(
+            **kwargs,
+        )
+
+
+@pytest.fixture()
 def sagemaker_session():
-    return Mock()
+    boto_mock = Mock(name="boto_session", region_name=REGION)
+    sms = MagicMock(
+        name="sagemaker_session",
+        boto_session=boto_mock,
+        boto_region_name=REGION,
+        config=None,
+        local_mode=False,
+        s3_client=None,
+        s3_resource=None,
+    )
+    sms.default_bucket = Mock(name="default_bucket", return_value=BUCKET_NAME)
+
+    return sms
 
 
 def test_prepare_container_def_with_model_data():
@@ -345,3 +375,75 @@ def test_delete_model_no_name(sagemaker_session):
     ):
         model.delete_model()
     sagemaker_session.delete_model.assert_not_called()
+
+
+@patch("time.strftime", MagicMock(return_value=TIMESTAMP))
+@patch("sagemaker.utils.repack_model")
+def test_script_mode_model_same_calls_as_framework(repack_model, sagemaker_session):
+    t = Model(
+        entry_point=ENTRY_POINT_INFERENCE,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        source_dir=SCRIPT_URI,
+        image_uri=IMAGE_URI,
+        model_data=MODEL_DATA,
+    )
+    t.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=INSTANCE_COUNT)
+
+    assert len(sagemaker_session.create_model.call_args_list) == 1
+    assert len(sagemaker_session.endpoint_from_production_variants.call_args_list) == 1
+    assert len(repack_model.call_args_list) == 1
+
+    generic_model_create_model_args = sagemaker_session.create_model.call_args_list
+    generic_model_endpoint_from_production_variants_args = (
+        sagemaker_session.endpoint_from_production_variants.call_args_list
+    )
+    generic_model_repack_model_args = repack_model.call_args_list
+
+    sagemaker_session.create_model.reset_mock()
+    sagemaker_session.endpoint_from_production_variants.reset_mock()
+    repack_model.reset_mock()
+
+    t = DummyFrameworkModel(
+        entry_point=ENTRY_POINT_INFERENCE,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        source_dir=SCRIPT_URI,
+        image_uri=IMAGE_URI,
+        model_data=MODEL_DATA,
+    )
+    t.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=INSTANCE_COUNT)
+
+    assert generic_model_create_model_args == sagemaker_session.create_model.call_args_list
+    assert (
+        generic_model_endpoint_from_production_variants_args
+        == sagemaker_session.endpoint_from_production_variants.call_args_list
+    )
+    assert generic_model_repack_model_args == repack_model.call_args_list
+
+
+@patch("sagemaker.git_utils.git_clone_repo")
+@patch("sagemaker.model.fw_utils.tar_and_upload_dir")
+def test_git_support_succeed_model_class(tar_and_upload_dir, git_clone_repo, sagemaker_session):
+    git_clone_repo.side_effect = lambda gitconfig, entrypoint, sourcedir, dependency: {
+        "entry_point": "entry_point",
+        "source_dir": "/tmp/repo_dir/source_dir",
+        "dependencies": ["/tmp/repo_dir/foo", "/tmp/repo_dir/bar"],
+    }
+    entry_point = "entry_point"
+    source_dir = "source_dir"
+    dependencies = ["foo", "bar"]
+    git_config = {"repo": GIT_REPO, "branch": BRANCH, "commit": COMMIT}
+    model = Model(
+        sagemaker_session=sagemaker_session,
+        entry_point=entry_point,
+        source_dir=source_dir,
+        dependencies=dependencies,
+        git_config=git_config,
+        image_uri=IMAGE_URI,
+    )
+    model.prepare_container_def(instance_type=INSTANCE_TYPE)
+    git_clone_repo.assert_called_with(git_config, entry_point, source_dir, dependencies)
+    assert model.entry_point == "entry_point"
+    assert model.source_dir == "/tmp/repo_dir/source_dir"
+    assert model.dependencies == ["/tmp/repo_dir/foo", "/tmp/repo_dir/bar"]