From 12e81dff69e93421d2f5b66d0174a3c049a7f629 Mon Sep 17 00:00:00 2001 From: Payton Staub Date: Thu, 29 Jul 2021 11:38:20 -0700 Subject: [PATCH] Use correct boto model name for RegisterModelStep properties --- src/sagemaker/workflow/_utils.py | 6 ++---- tests/unit/sagemaker/workflow/test_properties.py | 11 +++++++++++ 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/src/sagemaker/workflow/_utils.py b/src/sagemaker/workflow/_utils.py index 1abf53010e..0bff096651 100644 --- a/src/sagemaker/workflow/_utils.py +++ b/src/sagemaker/workflow/_utils.py @@ -284,9 +284,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: @@ -362,5 +360,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" + }