fix: Improve Pipeline workflow unit test branch coverage (#2878)

Co-authored-by: Dewen Qi <[email protected]>

Author: 2 people

src/sagemaker/workflow/_repack_model.py

@@ -34,7 +34,7 @@
 from distutils.dir_util import copy_tree
 
 
-def repack(inference_script, model_archive, dependencies=None, source_dir=None):
+def repack(inference_script, model_archive, dependencies=None, source_dir=None):  # pragma: no cover
     """Repack custom dependencies and code into an existing model TAR archive
 
     Args:
@@ -95,7 +95,7 @@ def repack(inference_script, model_archive, dependencies=None, source_dir=None):
         copy_tree(src_dir, "/opt/ml/model")
 
 
-if __name__ == "__main__":
+if __name__ == "__main__":  # pragma: no cover
     parser = argparse.ArgumentParser()
     parser.add_argument("--inference_script", type=str, default="inference.py")
     parser.add_argument("--dependencies", type=str, default=None)

src/sagemaker/workflow/_utils.py

@@ -80,7 +80,7 @@ def __init__(
                     artifacts. After the endpoint is created, the inference code
                     might use the IAM role, if it needs to access an AWS resource.
             model_data (str): The S3 location of a SageMaker model data
-                    ``.tar.gz`` file (default: None).
+                    ``.tar.gz`` file.
             entry_point (str): Path (absolute or relative) to the local Python
                     source file which should be executed as the entry point to
                     inference. If ``source_dir`` is specified, then ``entry_point``

src/sagemaker/workflow/condition_step.py

@@ -95,7 +95,7 @@ def properties(self):
 
 
 @attr.s
-class JsonGet(Expression):
+class JsonGet(Expression):  # pragma: no cover
     """Get JSON properties from PropertyFiles.
 
     Attributes:

src/sagemaker/workflow/functions.py

@@ -75,8 +75,8 @@ class JsonGet(Expression):
     @property
     def expr(self):
         """The expression dict for a `JsonGet` function."""
-        if not isinstance(self.step_name, str):
-            raise ValueError("Please give step name as a string")
+        if not isinstance(self.step_name, str) or not self.step_name:
+            raise ValueError("Please give a valid step name as a string")
 
         if isinstance(self.property_file, PropertyFile):
             name = self.property_file.name

src/sagemaker/workflow/lambda_step.py

@@ -161,8 +161,8 @@ def _get_function_arn(self):
             partition = "aws"
 
         if self.lambda_func.function_arn is None:
+            account_id = self.lambda_func.session.account_id()
             try:
-                account_id = self.lambda_func.session.account_id()
                 response = self.lambda_func.create()
                 return response["FunctionArn"]
             except ValueError as error:

src/sagemaker/workflow/step_collections.py

@@ -323,15 +323,15 @@ def __init__(
         """
         steps = []
         if "entry_point" in kwargs:
-            entry_point = kwargs["entry_point"]
-            source_dir = kwargs.get("source_dir")
-            dependencies = kwargs.get("dependencies")
+            entry_point = kwargs.get("entry_point", None)
+            source_dir = kwargs.get("source_dir", None)
+            dependencies = kwargs.get("dependencies", None)
             repack_model_step = _RepackModelStep(
                 name=f"{name}RepackModel",
                 depends_on=depends_on,
                 retry_policies=repack_model_step_retry_policies,
                 sagemaker_session=estimator.sagemaker_session,
-                role=estimator.sagemaker_session,
+                role=estimator.role,
                 model_data=model_data,
                 entry_point=entry_point,
                 source_dir=source_dir,
@@ -357,7 +357,11 @@ def predict_wrapper(endpoint, session):
             vpc_config=None,
             sagemaker_session=estimator.sagemaker_session,
             role=estimator.role,
-            **kwargs,
+            env=kwargs.get("env", None),
+            name=kwargs.get("name", None),
+            enable_network_isolation=kwargs.get("enable_network_isolation", None),
+            model_kms_key=kwargs.get("model_kms_key", None),
+            image_config=kwargs.get("image_config", None),
         )
         model_step = CreateModelStep(
             name=f"{name}CreateModelStep",

tests/integ/test_workflow.py

@@ -67,7 +67,6 @@
     ConditionLessThanOrEqualTo,
 )
 from sagemaker.workflow.condition_step import ConditionStep
-from sagemaker.workflow.condition_step import JsonGet as ConditionStepJsonGet
 from sagemaker.workflow.callback_step import (
     CallbackStep,
     CallbackOutput,
@@ -2835,8 +2834,8 @@ def test_end_to_end_pipeline_successful_execution(
 
     # define condition step
     cond_lte = ConditionLessThanOrEqualTo(
-        left=ConditionStepJsonGet(
-            step=step_eval,
+        left=JsonGet(
+            step_name=step_eval.name,
             property_file=evaluation_report,
             json_path="regression_metrics.mse.value",
         ),

tests/integ/test_workflow_with_clarify.py

@@ -33,7 +33,8 @@
 from sagemaker.processing import ProcessingInput, ProcessingOutput
 from sagemaker.session import get_execution_role
 from sagemaker.workflow.conditions import ConditionLessThanOrEqualTo
-from sagemaker.workflow.condition_step import ConditionStep, JsonGet
+from sagemaker.workflow.condition_step import ConditionStep
+from sagemaker.workflow.functions import JsonGet
 from sagemaker.workflow.parameters import (
     ParameterInteger,
     ParameterString,
@@ -237,7 +238,7 @@ def test_workflow_with_clarify(
         )
 
         cond_left = JsonGet(
-            step=step_process,
+            step_name=step_process.name,
             property_file="BiasOutput",
             json_path="post_training_bias_metrics.facets.F1[0].metrics[0].value",
         )

tests/unit/test_airflow.py renamed to tests/unit/sagemaker/workflow/test_airflow.py

@@ -13,6 +13,8 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import pytest
+
 from sagemaker.workflow.execution_variables import ExecutionVariables
 from sagemaker.workflow.functions import Join, JsonGet
 from sagemaker.workflow.parameters import (
@@ -97,3 +99,23 @@ def test_json_get_expressions():
             "Path": "my-json-path",
         },
     }
+
+
+def test_json_get_expressions_with_invalid_step_name():
+    with pytest.raises(ValueError) as err:
+        JsonGet(
+            step_name="",
+            property_file="my-property-file",
+            json_path="my-json-path",
+        ).expr
+
+    assert "Please give a valid step name as a string" in str(err.value)
+
+    with pytest.raises(ValueError) as err:
+        JsonGet(
+            step_name=ParameterString(name="MyString"),
+            property_file="my-property-file",
+            json_path="my-json-path",
+        ).expr
+
+    assert "Please give a valid step name as a string" in str(err.value)