fix: address review comments (iteration #2)

Author: aviruthen

sagemaker-train/tests/unit/train/test_model_trainer_pipeline_variable.py

@@ -182,7 +182,13 @@ class TestModelTrainerPipelineVariableHyperparameters:
     """Test that PipelineVariable objects work correctly in ModelTrainer hyperparameters."""
 
     def test_hyperparameters_with_parameter_integer(self):
-        """ParameterInteger in hyperparameters should be preserved through _create_training_job_args."""
+        """ParameterInteger in hyperparameters should be preserved through _create_training_job_args.
+
+        This test documents the exact bug scenario from GH#5504: safe_serialize
+        would fall back to str(data) for PipelineVariable objects, but
+        PipelineVariable.__str__ intentionally raises TypeError.
+        Before the fix, this call would have raised TypeError.
+        """
         max_depth = ParameterInteger(name="MaxDepth", default_value=5)
         trainer = ModelTrainer(
             training_image=DEFAULT_IMAGE,
@@ -192,6 +198,7 @@ def test_hyperparameters_with_parameter_integer(self):
             output_data_config=DEFAULT_OUTPUT,
             hyperparameters={"max_depth": max_depth},
         )
+        # This call would have raised TypeError before the fix (GH#5504)
         args = trainer._create_training_job_args()
         # PipelineVariable should be preserved as-is, not stringified
         assert args["hyper_parameters"]["max_depth"] is max_depth
@@ -210,32 +217,6 @@ def test_hyperparameters_with_parameter_string(self):
         args = trainer._create_training_job_args()
         assert args["hyper_parameters"]["algorithm"] is algo
 
-    def test_hyperparameters_with_parameter_integer_does_not_raise(self):
-        """Verify ParameterInteger in hyperparameters does NOT raise TypeError.
-
-        This test documents the exact bug scenario from GH#5504: safe_serialize
-        would fall back to str(data) for PipelineVariable objects, but
-        PipelineVariable.__str__ intentionally raises TypeError.
-        """
-        max_depth = ParameterInteger(name="MaxDepth", default_value=5)
-        trainer = ModelTrainer(
-            training_image=DEFAULT_IMAGE,
-            role=DEFAULT_ROLE,
-            compute=DEFAULT_COMPUTE,
-            stopping_condition=DEFAULT_STOPPING,
-            output_data_config=DEFAULT_OUTPUT,
-            hyperparameters={"max_depth": max_depth},
-        )
-        # This call would have raised TypeError before the fix
-        try:
-            args = trainer._create_training_job_args()
-        except TypeError:
-            pytest.fail(
-                "safe_serialize raised TypeError on PipelineVariable - "
-                "this is the bug described in GH#5504"
-            )
-        assert args["hyper_parameters"]["max_depth"] is max_depth
-
     def test_hyperparameters_with_mixed_pipeline_and_static_values(self):
         """Mixed PipelineVariable and static values should both be handled correctly."""
         max_depth = ParameterInteger(name="MaxDepth", default_value=5)

sagemaker-train/tests/unit/train/test_safe_serialize.py

@@ -31,27 +31,16 @@
 class TestSafeSerializeWithPipelineVariables:
     """Test safe_serialize handles PipelineVariable objects correctly."""
 
-    def test_safe_serialize_with_parameter_integer(self):
-        """ParameterInteger should be returned as-is (identity preserved)."""
-        param = ParameterInteger(name="MaxDepth", default_value=5)
+    @pytest.mark.parametrize("param", [
+        ParameterInteger(name="MaxDepth", default_value=5),
+        ParameterString(name="Algorithm", default_value="xgboost"),
+    ])
+    def test_safe_serialize_returns_pipeline_variable_as_is(self, param):
+        """PipelineVariable objects should be returned as-is (identity preserved)."""
         result = safe_serialize(param)
         assert result is param
         assert isinstance(result, PipelineVariable)
 
-    def test_safe_serialize_with_parameter_string(self):
-        """ParameterString should be returned as-is (identity preserved)."""
-        param = ParameterString(name="Algorithm", default_value="xgboost")
-        result = safe_serialize(param)
-        assert result is param
-        assert isinstance(result, PipelineVariable)
-
-    def test_safe_serialize_does_not_call_str_on_pipeline_variable(self):
-        """Verify that PipelineVariable.__str__ is never invoked (would raise TypeError)."""
-        param = ParameterInteger(name="TestParam", default_value=10)
-        # This should NOT raise TypeError
-        result = safe_serialize(param)
-        assert result is param
-
     def test_pipeline_variable_str_raises_type_error(self):
         """Confirm PipelineVariable.__str__ raises TypeError (the root cause of the bug)."""
         param = ParameterInteger(name="TestParam", default_value=10)