diff --git a/src/sagemaker/workflow/_utils.py b/src/sagemaker/workflow/_utils.py index ff967d4370..d32df19dfe 100644 --- a/src/sagemaker/workflow/_utils.py +++ b/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 diff --git a/tests/unit/sagemaker/workflow/test_properties.py b/tests/unit/sagemaker/workflow/test_properties.py index ff51d78346..5264a304ba 100644 --- a/tests/unit/sagemaker/workflow/test_properties.py +++ b/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" + }