diff --git a/expected_workflows/flow_hybrid_expected.py b/expected_workflows/flow_hybrid_expected.py
index 5c92307..c45d802 100644
--- a/expected_workflows/flow_hybrid_expected.py
+++ b/expected_workflows/flow_hybrid_expected.py
@@ -6,6 +6,7 @@ import json
 import datetime
 import re
 import jmespath
+from temporalio import workflow
 
 # Configure logging
 logging.basicConfig(level=logging.INFO,
@@ -16,7 +17,7 @@ logger = logging.getLogger(__name__)
 class test_repo_test_branch:
     @temporalio.workflow.run
     async def run(self, root_inputs: Dict[str, Any]) -> Dict[str, Any]:
-        await temporalio.workflow.sleep(0)
+        await workflow.sleep(0)
         workflow_info = temporalio.workflow.info()
         workflow_output: Dict[str, Any] = {
             "workflow_id": workflow_info.workflow_id,
diff --git a/expected_workflows/flow_parallel_expected.py b/expected_workflows/flow_parallel_expected.py
index aba06e4..1c9f407 100644
--- a/expected_workflows/flow_parallel_expected.py
+++ b/expected_workflows/flow_parallel_expected.py
@@ -6,6 +6,7 @@ import json
 import datetime
 import re
 import jmespath
+from temporalio import workflow
 
 # Configure logging
 logging.basicConfig(level=logging.INFO,
@@ -16,7 +17,7 @@ logger = logging.getLogger(__name__)
 class test_repo_test_branch:
     @temporalio.workflow.run
     async def run(self, root_inputs: Dict[str, Any]) -> Dict[str, Any]:
-        await temporalio.workflow.sleep(0)
+        await workflow.sleep(0)
         workflow_info = temporalio.workflow.info()
         workflow_output: Dict[str, Any] = {
             "workflow_id": workflow_info.workflow_id,
diff --git a/expected_workflows/flow_sequential_expected.py b/expected_workflows/flow_sequential_expected.py
index d0611ad..128fcba 100644
--- a/expected_workflows/flow_sequential_expected.py
+++ b/expected_workflows/flow_sequential_expected.py
@@ -6,6 +6,7 @@ import json
 import datetime
 import re
 import jmespath
+from temporalio import workflow
 
 # Configure logging
 logging.basicConfig(level=logging.INFO,
@@ -16,7 +17,7 @@ logger = logging.getLogger(__name__)
 class test_repo_test_branch:
     @temporalio.workflow.run
     async def run(self, root_inputs: Dict[str, Any]) -> Dict[str, Any]:
-        await temporalio.workflow.sleep(0)
+        await workflow.sleep(0)
         workflow_info = temporalio.workflow.info()
         workflow_output: Dict[str, Any] = {
             "workflow_id": workflow_info.workflow_id,