From 1ca6f77ef5542e63caf7fa470094780231b8713f Mon Sep 17 00:00:00 2001 From: aviruthen <91846056+aviruthen@users.noreply.github.com> Date: Tue, 7 Apr 2026 14:14:41 -0400 Subject: [PATCH] fix: Invalid import (5637) --- sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py | 2 +- sagemaker-mlops/tests/unit/workflow/test_steps.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py index 5e0eb3dda3..76e90a5309 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py @@ -205,7 +205,7 @@ def _find_dependencies_in_step_arguments( else: dependencies.add(self._get_step_name_from_str(referenced_step, step_map)) - from sagemaker.core.workflow.function_step import DelayedReturn + from sagemaker.mlops.workflow.function_step import DelayedReturn # TODO: we can remove the if-elif once move the validators to JsonGet constructor if isinstance(pipeline_variable, JsonGet): diff --git a/sagemaker-mlops/tests/unit/workflow/test_steps.py b/sagemaker-mlops/tests/unit/workflow/test_steps.py index 06cc729ceb..61bedcbc1a 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_steps.py +++ b/sagemaker-mlops/tests/unit/workflow/test_steps.py @@ -409,7 +409,7 @@ def test_step_find_dependencies_in_step_arguments_with_json_get(): obj = {"key": json_get} with patch('sagemaker.mlops.workflow.steps.TYPE_CHECKING', False): - with patch.dict('sys.modules', {'sagemaker.core.workflow.function_step': Mock()}): + with patch.dict('sys.modules', {'sagemaker.mlops.workflow.function_step': Mock()}): dependencies = Step._find_dependencies_in_step_arguments(step2, obj, {"step1": step1}) assert "step1" in dependencies @@ -445,7 +445,7 @@ def test_step_find_dependencies_in_step_arguments_with_delayed_return(): mock_module = Mock() mock_module.DelayedReturn = delayed_return_class - with patch.dict('sys.modules', {'sagemaker.core.workflow.function_step': mock_module}): + with patch.dict('sys.modules', {'sagemaker.mlops.workflow.function_step': mock_module}): dependencies = Step._find_dependencies_in_step_arguments(step2, obj, {"step1": step1}) assert "step1" in dependencies @@ -473,7 +473,7 @@ def test_step_find_dependencies_in_step_arguments_with_string_reference(): mock_module = Mock() mock_module.DelayedReturn = delayed_return_class - with patch.dict('sys.modules', {'sagemaker.core.workflow.function_step': mock_module}): + with patch.dict('sys.modules', {'sagemaker.mlops.workflow.function_step': mock_module}): dependencies = Step._find_dependencies_in_step_arguments(step2, obj, step_map) assert "step1" in dependencies