From 3f2ccd91cc6e8c0dbc4445ae44e697c031b98041 Mon Sep 17 00:00:00 2001
From: aviruthen <91846056+aviruthen@users.noreply.github.com>
Date: Fri, 27 Mar 2026 16:30:00 -0400
Subject: [PATCH] fix: Invalid import (5637)

---
 .../src/sagemaker/mlops/workflow/steps.py           |  2 +-
 sagemaker-mlops/tests/unit/workflow/test_steps.py   | 13 ++++++++++---
 2 files changed, 11 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..d03bc449ed 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,11 +473,18 @@ 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
 
 
+def test_delayed_return_import_from_correct_module():
+    """Verify that DelayedReturn can be imported from sagemaker.mlops.workflow.function_step."""
+    from sagemaker.mlops.workflow.function_step import DelayedReturn
+    assert DelayedReturn is not None
+    assert hasattr(DelayedReturn, '_to_json_get')
+
+
 def test_tuning_step_requires_step_args():
     from sagemaker.mlops.workflow.steps import TuningStep