diff --git a/sagemaker-train/tests/unit/train/test_dpo_trainer.py b/sagemaker-train/tests/unit/train/test_dpo_trainer.py
index 46c181e3ab..5dfae85bfd 100644
--- a/sagemaker-train/tests/unit/train/test_dpo_trainer.py
+++ b/sagemaker-train/tests/unit/train/test_dpo_trainer.py
@@ -549,7 +549,7 @@ def test_train_passes_sequence_length_to_serverless_config(self, mock_training_j
             mock_get_role, mock_get_options, mock_resolve_model):
         mock_validate_group.return_value = "test-group"
         mock_resolve_model.return_value = ("test-model", "test-model")
-        mock_get_sagemaker_session.return_value = Mock()
+        mock_get_sagemaker_session.return_value = Mock(sagemaker_config={})
         mock_fine_tuning_options = Mock()
         mock_fine_tuning_options.to_dict.return_value = {}
         mock_get_options.return_value = (mock_fine_tuning_options, "model-arn", False)
diff --git a/sagemaker-train/tests/unit/train/test_rlvr_trainer.py b/sagemaker-train/tests/unit/train/test_rlvr_trainer.py
index a273e5a470..3929b3cfa5 100644
--- a/sagemaker-train/tests/unit/train/test_rlvr_trainer.py
+++ b/sagemaker-train/tests/unit/train/test_rlvr_trainer.py
@@ -552,7 +552,7 @@ def test_train_passes_sequence_length_to_serverless_config(self, mock_training_j
             mock_get_role, mock_get_options, mock_resolve_model):
         mock_validate_group.return_value = "test-group"
         mock_resolve_model.return_value = ("test-model", "test-model")
-        mock_get_sagemaker_session.return_value = Mock()
+        mock_get_sagemaker_session.return_value = Mock(sagemaker_config={})
         mock_fine_tuning_options = Mock()
         mock_fine_tuning_options.to_dict.return_value = {}
         mock_get_options.return_value = (mock_fine_tuning_options, "model-arn", False)
diff --git a/sagemaker-train/tests/unit/train/test_sft_trainer.py b/sagemaker-train/tests/unit/train/test_sft_trainer.py
index 8bffa025de..00996c7858 100644
--- a/sagemaker-train/tests/unit/train/test_sft_trainer.py
+++ b/sagemaker-train/tests/unit/train/test_sft_trainer.py
@@ -563,7 +563,7 @@ def test_train_passes_sequence_length_to_serverless_config(self, mock_training_j
             mock_get_role, mock_get_options, mock_resolve_model):
         mock_validate_group.return_value = "test-group"
         mock_resolve_model.return_value = ("test-model", "test-model")
-        mock_get_sagemaker_session.return_value = Mock()
+        mock_get_sagemaker_session.return_value = Mock(sagemaker_config={})
         mock_fine_tuning_options = Mock()
         mock_fine_tuning_options.to_dict.return_value = {}
         mock_get_options.return_value = (mock_fine_tuning_options, "model-arn", False)