change: Add more tests

Author: trungleduc

tests/unit/sagemaker/feature_store/feature_processor/test_config_uploader.py

@@ -50,6 +50,10 @@ def runtime_env_manager():
     return mocked_runtime_env_manager
 
 
+def custom_file_filter():
+    pass
+
+
 @pytest.fixture
 def remote_decorator_config(sagemaker_session):
     return Mock(
@@ -73,6 +77,24 @@ def config_uploader(remote_decorator_config, runtime_env_manager):
     return ConfigUploader(remote_decorator_config, runtime_env_manager)
 
 
+@pytest.fixture
+def remote_decorator_config_with_filter(sagemaker_session):
+    return Mock(
+        _JobSettings,
+        sagemaker_session=sagemaker_session,
+        s3_root_uri="some_s3_uri",
+        s3_kms_key="some_kms",
+        spark_config=SparkConfig(),
+        dependencies=None,
+        include_local_workdir=True,
+        pre_execution_commands="some_commands",
+        pre_execution_script="some_path",
+        python_sdk_whl_s3_uri=SAGEMAKER_WHL_FILE_S3_PATH,
+        environment_variables={"REMOTE_FUNCTION_SECRET_KEY": "some_secret_key"},
+        custom_file_filter=custom_file_filter,
+    )
+
+
 @patch("sagemaker.feature_store.feature_processor._config_uploader.StoredFunction")
 def test_prepare_and_upload_callable(mock_stored_function, config_uploader, wrapped_func):
     mock_stored_function.save(wrapped_func).return_value = None
@@ -113,6 +135,41 @@ def test_prepare_and_upload_dependencies(mock_upload, config_uploader):
     )
 
 
+@patch(
+    "sagemaker.feature_store.feature_processor._config_uploader._prepare_and_upload_dependencies",
+    return_value="some_s3_uri",
+)
+def test_prepare_and_upload_dependencies_with_filter(
+    mock_job_upload, remote_decorator_config_with_filter, runtime_env_manager
+):
+    config_uploader_with_filter = ConfigUploader(
+        remote_decorator_config=remote_decorator_config_with_filter,
+        runtime_env_manager=runtime_env_manager,
+    )
+    remote_decorator_config = config_uploader_with_filter.remote_decorator_config
+    config_uploader_with_filter._prepare_and_upload_dependencies(
+        local_dependencies_path="some/path/to/dependency",
+        include_local_workdir=True,
+        pre_execution_commands=remote_decorator_config.pre_execution_commands,
+        pre_execution_script_local_path=remote_decorator_config.pre_execution_script,
+        s3_base_uri=remote_decorator_config.s3_root_uri,
+        s3_kms_key=remote_decorator_config.s3_kms_key,
+        sagemaker_session=sagemaker_session,
+        custom_file_filter=remote_decorator_config_with_filter.custom_file_filter,
+    )
+
+    mock_job_upload.assert_called_once_with(
+        local_dependencies_path="some/path/to/dependency",
+        include_local_workdir=True,
+        pre_execution_commands=remote_decorator_config.pre_execution_commands,
+        pre_execution_script_local_path=remote_decorator_config.pre_execution_script,
+        s3_base_uri=remote_decorator_config.s3_root_uri,
+        s3_kms_key=remote_decorator_config.s3_kms_key,
+        sagemaker_session=sagemaker_session,
+        custom_file_filter=custom_file_filter,
+    )
+
+
 @patch(
     "sagemaker.feature_store.feature_processor._config_uploader._prepare_and_upload_runtime_scripts",
     return_value="some_s3_uri",

tests/unit/sagemaker/remote_function/test_client.py

@@ -175,6 +175,32 @@ def square(x):
     assert mock_job_settings.call_args.kwargs["image_uri"] == IMAGE
 
 
+@patch(
+    "sagemaker.remote_function.client.serialization.deserialize_obj_from_s3",
+    return_value=EXPECTED_JOB_RESULT,
+)
+@patch("sagemaker.remote_function.client._JobSettings")
+@patch("sagemaker.remote_function.client._Job.start")
+def test_decorator_with_custom_file_filter(
+    mock_start, mock_job_settings, mock_deserialize_obj_from_s3
+):
+    mock_job = Mock(job_name=TRAINING_JOB_NAME)
+    mock_job.describe.return_value = COMPLETED_TRAINING_JOB
+
+    mock_start.return_value = mock_job
+
+    def custom_file_filter():
+        pass
+
+    @remote(image_uri=IMAGE, s3_root_uri=S3_URI, custom_file_filter=custom_file_filter)
+    def square(x):
+        return x * x
+
+    result = square(5)
+    assert result == EXPECTED_JOB_RESULT
+    assert mock_job_settings.call_args.kwargs["custom_file_filter"] == custom_file_filter
+
+
 @patch(
     "sagemaker.remote_function.client.serialization.deserialize_exception_from_s3",
     return_value=ZeroDivisionError("division by zero"),