Add unit tests for ModelTrainer (#1527)

* Add unit tests for ModelTrainer

* Flake8

* format

Author: benieric

src/sagemaker/modules/train/model_trainer.py

@@ -108,7 +108,7 @@ class ModelTrainer(BaseModel):
             The VPC configuration.
     """
 
-    model_config = ConfigDict(arbitrary_types_allowed=True)
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
 
     session: Optional[Session] = None
     role: Optional[str] = None
@@ -274,19 +274,20 @@ def train(
         if input_data_channels:
             self.input_data_channels = input_data_channels
         if source_code_config:
+            self._validate_source_code_config(source_code_config)
             self.source_code_config = source_code_config
         if hyper_parameters:
             self.hyper_parameters = hyper_parameters
         if environment:
             self.environment = environment
 
-        input_data_config = self._get_input_data_config(self.input_data_channels)
+        input_data_config = []
+        if self.input_data_channels:
+            input_data_config = self._get_input_data_config(self.input_data_channels)
 
         container_entrypoint = None
         container_arguments = None
         if self.source_code_config:
-            if not input_data_config:
-                input_data_config = []
 
             # If source code is provided, create a channel for the source code
             # The source code will be mounted at /opt/ml/input/data/code in the container
@@ -397,9 +398,14 @@ def create_input_data_channel(self, channel_name: str, data_source: DataSourceTy
             else:
                 raise ValueError(f"Not a valid S3 URI or local file path: {data_source}.")
         elif isinstance(data_source, S3DataSource):
-            channel = Channel(channel_name=channel_name, data_source=data_source)
+            channel = Channel(
+                channel_name=channel_name, data_source=DataSource(s3_data_source=data_source)
+            )
         elif isinstance(data_source, FileSystemDataSource):
-            channel = Channel(channel_name=channel_name, data_source=data_source)
+            channel = Channel(
+                channel_name=channel_name,
+                data_source=DataSource(file_system_data_source=data_source),
+            )
         return channel
 
     def _get_input_data_config(

src/sagemaker/modules/utils.py

@@ -19,19 +19,6 @@
 from typing import Literal
 
 
-def _is_valid_ecr_image(image: str) -> bool:
-    """Check if the image is a valid ECR image URI.
-
-    Args:
-        image (str): The image URI sting to validate
-
-    Returns:
-        bool: True if the image is a valid ECR image URI, False otherwise
-    """
-    pattern = r"^\d{12}\.dkr\.ecr\.\w+-\d{1,2}\.amazonaws\.com\/[\w-]+(:[\w.-]+)?$"
-    return bool(re.match(pattern, image))
-
-
 def _is_valid_s3_uri(path: str, path_type: Literal["File", "Directory", "Any"] = "Any") -> bool:
     """Check if the path is a valid S3 URI.
 

tests/unit/sagemaker/modules/test_utils.py

@@ -0,0 +1,140 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Utils Tests."""
+from __future__ import absolute_import
+
+import pytest
+
+from tests.unit import DATA_DIR
+from sagemaker.modules.utils import (
+    _is_valid_s3_uri,
+    _is_valid_path,
+    _get_unique_name,
+    _get_repo_name_from_image,
+)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        {
+            "path": "s3://bucket/key",
+            "path_type": "Any",
+            "expected": True,
+        },
+        {
+            "path": "s3://bucket/key",
+            "path_type": "File",
+            "expected": True,
+        },
+        {
+            "path": "s3://bucket/key/",
+            "path_type": "Directory",
+            "expected": True,
+        },
+        {
+            "path": "s3://bucket/key/",
+            "path_type": "File",
+            "expected": False,
+        },
+        {
+            "path": "s3://bucket/key",
+            "path_type": "Directory",
+            "expected": False,
+        },
+        {
+            "path": "/bucket/key",
+            "path_type": "Any",
+            "expected": False,
+        },
+    ],
+)
+def test_is_valid_s3_uri(test_case):
+    assert _is_valid_s3_uri(test_case["path"], test_case["path_type"]) == test_case["expected"]
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        {
+            "path": DATA_DIR,
+            "path_type": "Any",
+            "expected": True,
+        },
+        {
+            "path": DATA_DIR,
+            "path_type": "Directory",
+            "expected": True,
+        },
+        {
+            "path": f"{DATA_DIR}/dummy_input.txt",
+            "path_type": "File",
+            "expected": True,
+        },
+        {
+            "path": f"{DATA_DIR}/dummy_input.txt",
+            "path_type": "Directory",
+            "expected": False,
+        },
+        {
+            "path": f"{DATA_DIR}/non_existent",
+            "path_type": "Any",
+            "expected": False,
+        },
+    ],
+)
+def test_is_valid_path(test_case):
+    assert _is_valid_path(test_case["path"], test_case["path_type"]) == test_case["expected"]
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        {
+            "base": "test",
+            "max_length": 5,
+        },
+        {
+            "base": "1111111111" * 7,
+            "max_length": None,
+        },
+    ],
+)
+def test_get_unique_name(test_case):
+    assert (
+        len(_get_unique_name(test_case["base"], test_case.get("max_length")))
+        <= test_case["max_length"]
+        if test_case.get("max_length")
+        else 63
+    )
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        {
+            "image": "123456789012.dkr.ecr.us-west-2.amazonaws.com/my-custom-image:latest",
+            "expected": "my-custom-image",
+        },
+        {
+            "image": "my-custom-image:latest",
+            "expected": "my-custom-image",
+        },
+        {
+            "image": "public.ecr.aws/docker/library/my-custom-image:latest",
+            "expected": "my-custom-image",
+        },
+    ],
+)
+def test_get_repo_name_from_image(test_case):
+    assert _get_repo_name_from_image(test_case["image"]) == test_case["expected"]