Fix unit-tests

tests/unit/sagemaker/huggingface/test_estimator.py

@@ -241,7 +241,7 @@ def test_huggingface(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         huggingface_training_version, f"pytorch{huggingface_pytorch_training_version}"

tests/unit/sagemaker/training_compiler/test_huggingface_pytorch_compiler.py

@@ -402,7 +402,7 @@ def test_pytorchxla_distribution(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         huggingface_training_compiler_version,
@@ -463,7 +463,7 @@ def test_default_compiler_config(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         huggingface_training_compiler_version,
@@ -519,7 +519,7 @@ def test_debug_compiler_config(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         huggingface_training_compiler_version,
@@ -575,7 +575,7 @@ def test_disable_compiler_config(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         huggingface_training_compiler_version,

tests/unit/sagemaker/training_compiler/test_huggingface_tensorflow_compiler.py

@@ -349,7 +349,7 @@ def test_default_compiler_config(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         huggingface_training_compiler_version,
@@ -407,7 +407,7 @@ def test_debug_compiler_config(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         huggingface_training_compiler_version,
@@ -465,7 +465,7 @@ def test_disable_compiler_config(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         huggingface_training_compiler_version,

tests/unit/sagemaker/training_compiler/test_pytorch_compiler.py

@@ -344,7 +344,7 @@ def test_pytorchxla_distribution(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         pytorch_training_compiler_version,
@@ -403,7 +403,7 @@ def test_default_compiler_config(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         pytorch_training_compiler_version,
@@ -458,7 +458,7 @@ def test_debug_compiler_config(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         pytorch_training_compiler_version,
@@ -513,7 +513,7 @@ def test_disable_compiler_config(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         pytorch_training_compiler_version,

tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py

@@ -289,7 +289,7 @@ def test_default(
         sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
         assert sagemaker_call_names == ["train", "logs_for_job"]
         boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-        assert boto_call_names == ["resource"]
+        assert boto_call_names == ["resource", "client"]
 
         expected_train_args = _create_train_job(
             tensorflow_training_version,
@@ -348,7 +348,7 @@ def test_byoc(
         sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
         assert sagemaker_call_names == ["train", "logs_for_job"]
         boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-        assert boto_call_names == ["resource"]
+        assert boto_call_names == ["resource", "client"]
 
         expected_train_args = _create_train_job(
             tensorflow_training_version,
@@ -399,7 +399,7 @@ def test_debug_compiler_config(
         sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
         assert sagemaker_call_names == ["train", "logs_for_job"]
         boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-        assert boto_call_names == ["resource"]
+        assert boto_call_names == ["resource", "client"]
 
         expected_train_args = _create_train_job(
             tensorflow_training_version,
@@ -450,7 +450,7 @@ def test_disable_compiler_config(
         sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
         assert sagemaker_call_names == ["train", "logs_for_job"]
         boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-        assert boto_call_names == ["resource"]
+        assert boto_call_names == ["resource", "client"]
 
         expected_train_args = _create_train_job(
             tensorflow_training_version,

tests/unit/test_chainer.py

@@ -354,7 +354,7 @@ def test_chainer(strftime, time, sagemaker_session, chainer_version, chainer_py_
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(chainer_version, chainer_py_version)
     expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs

tests/unit/test_mxnet.py

@@ -360,7 +360,7 @@ def test_mxnet(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     actual_train_args = sagemaker_session.method_calls[0][2]
     job_name = actual_train_args["job_name"]

tests/unit/test_processing.py

@@ -133,6 +133,14 @@ def uploaded_code(
     return UploadedCode(s3_prefix=s3_prefix, script_name=script_name)
 
 
+@pytest.fixture(auto_use=True)
+def mock_process_studio_metadata_file():
+    with patch(
+        "sagemaker.user_agent.process_studio_metadata_file", MagicMock()
+    ) as process_studio_metadata_file:
+        yield process_studio_metadata_file
+
+
 @patch("sagemaker.utils._botocore_resolver")
 @patch("os.path.exists", return_value=True)
 @patch("os.path.isfile", return_value=True)
@@ -277,15 +285,13 @@ def test_sklearn_with_all_parameters_via_run_args(
     sagemaker_session.process.assert_called_with(**expected_args)
 
 
-@patch("sagemaker.user_agent.process_studio_metadata_file", return_value="TestApp")
 @patch("sagemaker.utils._botocore_resolver")
 @patch("os.path.exists", return_value=True)
 @patch("os.path.isfile", return_value=True)
 def test_sklearn_with_all_parameters_via_run_args_called_twice(
     exists_mock,
     isfile_mock,
     botocore_resolver,
-    process_studio_metadata_file_mock,
     sklearn_version,
     sagemaker_session,
 ):
@@ -338,15 +344,13 @@ def test_sklearn_with_all_parameters_via_run_args_called_twice(
     sagemaker_session.process.assert_called_with(**expected_args)
 
 
-@patch("sagemaker.user_agent.process_studio_metadata_file", return_value="TestApp")
 @patch("sagemaker.utils._botocore_resolver")
 @patch("os.path.exists", return_value=True)
 @patch("os.path.isfile", return_value=True)
 def test_pytorch_processor_with_required_parameters(
     exists_mock,
     isfile_mock,
     botocore_resolver,
-    process_studio_metadata_file_mock,
     sagemaker_session,
     pytorch_training_version,
     pytorch_training_py_version,
@@ -418,14 +422,12 @@ def test_xgboost_processor_with_required_parameters(
     sagemaker_session.process.assert_called_with(**expected_args)
 
 
-@patch("sagemaker.user_agent.process_studio_metadata_file", return_value="TestApp")
 @patch("sagemaker.utils._botocore_resolver")
 @patch("os.path.exists", return_value=True)
 @patch("os.path.isfile", return_value=True)
 def test_mxnet_processor_with_required_parameters(
     exists_mock,
     isfile_mock,
-    process_studio_metadata_file_mock,
     botocore_resolver,
     sagemaker_session,
     mxnet_training_version,
@@ -468,15 +470,13 @@ def test_mxnet_processor_with_required_parameters(
     sagemaker_session.process.assert_called_with(**expected_args)
 
 
-@patch("sagemaker.user_agent.process_studio_metadata_file", return_value="TestApp")
 @patch("sagemaker.utils._botocore_resolver")
 @patch("os.path.exists", return_value=True)
 @patch("os.path.isfile", return_value=True)
 def test_tensorflow_processor_with_required_parameters(
     exists_mock,
     isfile_mock,
     botocore_resolver,
-    process_studio_metadata_file_mock,
     sagemaker_session,
     tensorflow_training_version,
     tensorflow_training_py_version,

tests/unit/test_pytorch.py

@@ -337,7 +337,7 @@ def test_pytorch(
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(pytorch_inference_version, pytorch_inference_py_version)
     expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs

tests/unit/test_rl.py

@@ -335,7 +335,7 @@ def test_rl(time, strftime, sagemaker_session, coach_mxnet_version):
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         RLToolkit.COACH.value, coach_mxnet_version, RLFramework.MXNET.value

tests/unit/test_sklearn.py

@@ -332,7 +332,7 @@ def test_sklearn(time, strftime, sagemaker_session, sklearn_version):
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(sklearn_version)
     expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs

tests/unit/test_xgboost.py

@@ -330,7 +330,7 @@ def test_xgboost_cpu(time, strftime, sagemaker_session, xgboost_framework_versio
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(xgboost_framework_version)
     expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs
@@ -377,7 +377,7 @@ def test_xgboost_gpu(time, strftime, sagemaker_session, xgboost_gpu_framework_ve
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(
         xgboost_gpu_framework_version, instance_type=GPU_INSTANCE_TYPE
@@ -427,7 +427,7 @@ def test_distributed_training(time, strftime, sagemaker_session, xgboost_framewo
     sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls]
     assert sagemaker_call_names == ["train", "logs_for_job"]
     boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls]
-    assert boto_call_names == ["resource"]
+    assert boto_call_names == ["resource", "client"]
 
     expected_train_args = _create_train_job(xgboost_framework_version, DIST_INSTANCE_COUNT)
     expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs