Use correct boto model name for RegisterModelStep properties (#2543)

Co-authored-by: Payton Staub <[email protected]>
Co-authored-by: Shreya Pandit <[email protected]>
Co-authored-by: Ahsan Khan <[email protected]>

Author: 4 people

src/sagemaker/workflow/_utils.py

@@ -288,9 +288,7 @@ def __init__(
         self.kwargs = kwargs
         self.container_def_list = container_def_list
 
-        self._properties = Properties(
-            path=f"Steps.{name}", shape_name="DescribeModelPackageResponse"
-        )
+        self._properties = Properties(path=f"Steps.{name}", shape_name="DescribeModelPackageOutput")
 
     @property
     def arguments(self) -> RequestType:
@@ -366,5 +364,5 @@ def arguments(self) -> RequestType:
 
     @property
     def properties(self):
-        """A Properties object representing the DescribeTrainingJobResponse data model."""
+        """A Properties object representing the DescribeModelPackageOutput data model."""
         return self._properties

tests/unit/sagemaker/workflow/test_properties.py

@@ -67,3 +67,14 @@ def test_properties_tuning_job():
     assert prop.TrainingJobSummaries[0].TrainingJobName.expr == {
         "Get": "Steps.MyStep.TrainingJobSummaries[0].TrainingJobName"
     }
+
+
+def test_properties_describe_model_package_output():
+    prop = Properties("Steps.MyStep", "DescribeModelPackageOutput")
+    some_prop_names = ["ModelPackageName", "ModelPackageGroupName", "ModelPackageArn"]
+    for name in some_prop_names:
+        assert name in prop.__dict__.keys()
+    assert prop.ModelPackageName.expr == {"Get": "Steps.MyStep.ModelPackageName"}
+    assert prop.ValidationSpecification.ValidationRole.expr == {
+        "Get": "Steps.MyStep.ValidationSpecification.ValidationRole"
+    }