added unit test

Author: Basil Beirouti

tests/unit/sagemaker/model/test_model.py

@@ -48,6 +48,47 @@
 IMAGE_URI = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.9.0-gpu-py38"
 
 
+MODEL_DESCRIPTION = "a description"
+
+SUPPORTED_REALTIME_INFERENCE_INSTANCE_TYPES = ["ml.m4.xlarge"]
+SUPPORTED_BATCH_TRANSFORM_INSTANCE_TYPES = ["ml.m4.xlarge"]
+
+SUPPORTED_CONTENT_TYPES = ["text/csv", "application/json", "application/jsonlines"]
+SUPPORTED_RESPONSE_MIME_TYPES = ["application/json", "text/csv", "application/jsonlines"]
+
+VALIDATION_FILE_NAME = "input.csv"
+VALIDATION_INPUT_PATH = "s3://" + BUCKET_NAME + "/validation-input-csv/"
+VALIDATION_OUTPUT_PATH = "s3://" + BUCKET_NAME + "/validation-output-csv/"
+
+VALIDATION_SPECIFICATION = {
+    "ValidationRole": "some_role",
+    "ValidationProfiles": [
+        {
+            "ProfileName": "Validation-test",
+            "TransformJobDefinition": {
+                "BatchStrategy": "SingleRecord",
+                "TransformInput": {
+                    "DataSource": {
+                        "S3DataSource": {
+                            "S3DataType": "S3Prefix",
+                            "S3Uri": VALIDATION_INPUT_PATH,
+                        }
+                    },
+                    "ContentType": SUPPORTED_CONTENT_TYPES[0],
+                },
+                "TransformOutput": {
+                    "S3OutputPath": VALIDATION_OUTPUT_PATH,
+                },
+                "TransformResources": {
+                    "InstanceType": SUPPORTED_BATCH_TRANSFORM_INSTANCE_TYPES[0],
+                    "InstanceCount": 1,
+                },
+            },
+        },
+    ],
+}
+
+
 class DummyFrameworkModel(FrameworkModel):
     def __init__(self, **kwargs):
         super(DummyFrameworkModel, self).__init__(
@@ -687,3 +728,40 @@ def test_script_mode_model_uses_proper_sagemaker_submit_dir(repack_model, sagema
         ]
         == "/opt/ml/model/code"
     )
+
+
+@patch("sagemaker.get_model_package_args")
+def test_call_to_get_model_package_args(get_model_package_args, sagemaker_session):
+
+    source_dir = "s3://blah/blah/blah"
+    t = Model(
+        entry_point=ENTRY_POINT_INFERENCE,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        source_dir=source_dir,
+        image_uri=IMAGE_URI,
+        model_data=MODEL_DATA,
+    )
+
+    t.register(
+        SUPPORTED_CONTENT_TYPES,
+        SUPPORTED_RESPONSE_MIME_TYPES,
+        SUPPORTED_REALTIME_INFERENCE_INSTANCE_TYPES,
+        SUPPORTED_BATCH_TRANSFORM_INSTANCE_TYPES,
+        marketplace_cert=True,
+        description=MODEL_DESCRIPTION,
+        model_package_name=MODEL_NAME,
+        validation_specification=VALIDATION_SPECIFICATION,
+    )
+
+    # check that the kwarg validation_specification was passed to the internal method 'get_model_package_args'
+    assert (
+        "validation_specification" in get_model_package_args.call_args_list[0][1]
+    ), "validation_specification kwarg was not passed to get_model_package_args"
+
+    # check that the kwarg validation_specification is identical to the one passed into the method 'register'
+    assert (
+        VALIDATION_SPECIFICATION
+        == get_model_package_args.call_args_list[0][1]["validation_specification"]
+    ), """ValidationSpecification from model.register method is not identical to validation_spec from
+         get_model_package_args"""