another bunch of minor changes

Signed-off-by: Tim Li <[email protected]>

Author: timl3136

cadence/_internal/workflow/workflow_engine.py

@@ -20,9 +20,12 @@ class DecisionResult:
     decisions: list[Decision]
 
 class WorkflowEngine:
-    def __init__(self, info: WorkflowInfo, client: Client, workflow_func: Callable[[Any], Any] | None = None):
+    def __init__(self, info: WorkflowInfo, client: Client, workflow_definition=None):
         self._context = Context(client, info)
-        self._workflow_func = workflow_func
+        self._workflow_definition = workflow_definition
+        self._workflow_instance = None
+        if workflow_definition:
+            self._workflow_instance = workflow_definition.cls()
         self._decision_manager = DecisionManager()
         self._decisions_helper = DecisionsHelper(self._decision_manager)
         self._is_workflow_complete = False
@@ -250,19 +253,17 @@ def _fallback_process_workflow_history(self, history) -> None:
     async def _execute_workflow_function(self, decision_task: PollForDecisionTaskResponse) -> None:
         """
         Execute the workflow function to generate new decisions.
-        
+
         This blocks until the workflow schedules an activity or completes.
-        
+
         Args:
             decision_task: The decision task containing workflow context
         """
         try:
-            # Execute the workflow function
-            # The workflow function should block until it schedules an activity
-            workflow_func = self._workflow_func
-            if workflow_func is None:
+            # Execute the workflow function from the workflow instance
+            if self._workflow_definition is None or self._workflow_instance is None:
                 logger.warning(
-                    "No workflow function available",
+                    "No workflow definition or instance available",
                     extra={
                         "workflow_type": self._context.info().workflow_type,
                         "workflow_id": self._context.info().workflow_id,
@@ -271,6 +272,9 @@ async def _execute_workflow_function(self, decision_task: PollForDecisionTaskRes
                 )
                 return
 
+            # Get the workflow run method from the instance
+            workflow_func = self._workflow_definition.get_run_method(self._workflow_instance)
+
             # Extract workflow input from history
             workflow_input = await self._extract_workflow_input(decision_task)
 
@@ -290,7 +294,7 @@ async def _execute_workflow_function(self, decision_task: PollForDecisionTaskRes
                         "completion_type": "success"
                     }
                 )
-                
+
         except Exception as e:
             logger.error(
                 "Error executing workflow function",

cadence/worker/_decision_task_handler.py

@@ -103,9 +103,9 @@ async def _handle_task_implementation(self, task: PollForDecisionTaskResponse) -
             workflow_engine = self._workflow_engines.get(cache_key)
             if workflow_engine is None:
                 workflow_engine = WorkflowEngine(
-                    info=workflow_info, 
-                    client=self._client, 
-                    workflow_func=workflow_definition.fn
+                    info=workflow_info,
+                    client=self._client,
+                    workflow_definition=workflow_definition
                 )
                 self._workflow_engines[cache_key] = workflow_engine
 

cadence/workflow.py

@@ -2,7 +2,7 @@
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import Iterator, Callable, TypeVar, TypedDict, Type
+from typing import Iterator, Callable, TypeVar, TypedDict, Type, cast, Any
 from functools import wraps
 
 from cadence.client import Client
@@ -37,6 +37,16 @@ def cls(self) -> Type:
         """Get the workflow class."""
         return self._cls
 
+    def get_run_method(self, instance: Any) -> Callable:
+        """Get the workflow run method from an instance of the workflow class."""
+        for attr_name in dir(instance):
+            if attr_name.startswith('_'):
+                continue
+            attr = getattr(instance, attr_name)
+            if callable(attr) and hasattr(attr, '_workflow_run'):
+                return cast(Callable, attr)
+        raise ValueError(f"No @workflow.run method found in class {self._cls.__name__}")
+
     @staticmethod
     def wrap(cls: Type, opts: WorkflowDefinitionOptions) -> 'WorkflowDefinition':
         """

tests/cadence/_internal/workflow/test_workflow_engine_integration.py

@@ -36,19 +36,25 @@ def workflow_info(self):
         )
 
     @pytest.fixture
-    def mock_workflow_func(self):
-        """Create a mock workflow function."""
-        def workflow_func(input_data):
-            return f"processed: {input_data}"
-        return workflow_func
+    def mock_workflow_definition(self):
+        """Create a mock workflow definition."""
+        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions, workflow
+
+        class TestWorkflow:
+            @workflow.run
+            def weird_name(self, input_data):
+                return f"processed: {input_data}"
+
+        workflow_opts = WorkflowDefinitionOptions(name="test_workflow")
+        return WorkflowDefinition.wrap(TestWorkflow, workflow_opts)
 
     @pytest.fixture
-    def workflow_engine(self, mock_client, workflow_info, mock_workflow_func):
+    def workflow_engine(self, mock_client, workflow_info, mock_workflow_definition):
         """Create a WorkflowEngine instance."""
         return WorkflowEngine(
             info=workflow_info,
             client=mock_client,
-            workflow_func=mock_workflow_func
+            workflow_definition=mock_workflow_definition
         )
 
     def create_mock_decision_task(self, workflow_id="test-workflow", run_id="test-run", workflow_type="test_workflow"):
@@ -211,10 +217,13 @@ async def test_extract_workflow_input_deserialization_error(self, workflow_engin
     def test_execute_workflow_function_sync(self, workflow_engine):
         """Test synchronous workflow function execution."""
         input_data = "test-input"
-        
+
+        # Get the workflow function from the instance
+        workflow_func = workflow_engine._workflow_definition.get_run_method(workflow_engine._workflow_instance)
+
         # Execute the workflow function
-        result = workflow_engine._execute_workflow_function_once(workflow_engine._workflow_func, input_data)
-        
+        result = workflow_engine._execute_workflow_function_once(workflow_func, input_data)
+
         # Verify the result
         assert result == "processed: test-input"
 
@@ -239,20 +248,21 @@ def test_execute_workflow_function_none(self, workflow_engine):
         with pytest.raises(TypeError, match="'NoneType' object is not callable"):
             workflow_engine._execute_workflow_function_once(None, input_data)
 
-    def test_workflow_engine_initialization(self, workflow_engine, workflow_info, mock_client, mock_workflow_func):
+    def test_workflow_engine_initialization(self, workflow_engine, workflow_info, mock_client, mock_workflow_definition):
         """Test WorkflowEngine initialization."""
         assert workflow_engine._context is not None
-        assert workflow_engine._workflow_func == mock_workflow_func
+        assert workflow_engine._workflow_definition == mock_workflow_definition
+        assert workflow_engine._workflow_instance is not None
         assert workflow_engine._decision_manager is not None
         assert workflow_engine._is_workflow_complete is False
 
     @pytest.mark.asyncio
-    async def test_workflow_engine_without_workflow_func(self, mock_client, workflow_info):
-        """Test WorkflowEngine without workflow function."""
+    async def test_workflow_engine_without_workflow_definition(self, mock_client, workflow_info):
+        """Test WorkflowEngine without workflow definition."""
         engine = WorkflowEngine(
             info=workflow_info,
             client=mock_client,
-            workflow_func=None
+            workflow_definition=None
         )
 
         decision_task = self.create_mock_decision_task()
@@ -269,12 +279,21 @@ async def test_workflow_engine_without_workflow_func(self, mock_client, workflow
     async def test_workflow_engine_workflow_completion(self, workflow_engine, mock_client):
         """Test workflow completion detection."""
         decision_task = self.create_mock_decision_task()
-        
-        # Mock workflow function to return a result (indicating completion)
-        def completing_workflow_func(input_data):
-            return "workflow-completed"
-        
-        workflow_engine._workflow_func = completing_workflow_func
+
+        # Create a workflow definition that returns a result (indicating completion)
+        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions, workflow
+
+        class CompletingWorkflow:
+            @workflow.run
+            def run(self, input_data):
+                return "workflow-completed"
+
+        workflow_opts = WorkflowDefinitionOptions(name="completing_workflow")
+        completing_definition = WorkflowDefinition.wrap(CompletingWorkflow, workflow_opts)
+
+        # Replace the workflow definition and instance
+        workflow_engine._workflow_definition = completing_definition
+        workflow_engine._workflow_instance = completing_definition.cls()
 
         with patch.object(workflow_engine._decision_manager, 'collect_pending_decisions', return_value=[]):
             # Process the decision

tests/cadence/worker/test_decision_task_handler.py

@@ -83,12 +83,15 @@ def test_initialization(self, mock_client, mock_registry):
     async def test_handle_task_implementation_success(self, handler, sample_decision_task, mock_registry):
         """Test successful decision task handling."""
         # Create actual workflow definition
-        def mock_workflow_func():
-            return "test_result"
+        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions, workflow
+
+        class MockWorkflow:
+            @workflow.run
+            async def run(self):
+                return "test_result"
 
-        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions
         workflow_opts = WorkflowDefinitionOptions(name="test_workflow")
-        workflow_definition = WorkflowDefinition.wrap(mock_workflow_func, workflow_opts)
+        workflow_definition = WorkflowDefinition.wrap(MockWorkflow, workflow_opts)
         mock_registry.get_workflow.return_value = workflow_definition
 
         # Mock workflow engine
@@ -148,12 +151,15 @@ async def test_handle_task_implementation_workflow_not_found(self, handler, samp
     async def test_handle_task_implementation_caches_engines(self, handler, sample_decision_task, mock_registry):
         """Test that decision task handler caches workflow engines for same workflow execution."""
         # Create actual workflow definition
-        def mock_workflow_func():
-            return "test_result"
+        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions, workflow
+
+        class MockWorkflow:
+            @workflow.run
+            async def run(self):
+                return "test_result"
 
-        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions
         workflow_opts = WorkflowDefinitionOptions(name="test_workflow")
-        workflow_definition = WorkflowDefinition.wrap(mock_workflow_func, workflow_opts)
+        workflow_definition = WorkflowDefinition.wrap(MockWorkflow, workflow_opts)
         mock_registry.get_workflow.return_value = workflow_definition
 
         # Mock workflow engine
@@ -182,9 +188,17 @@ def mock_workflow_func():
     @pytest.mark.asyncio
     async def test_handle_task_implementation_different_executions_get_separate_engines(self, handler, mock_registry):
         """Test that different workflow executions get separate engines."""
-        # Mock workflow function
-        mock_workflow_func = Mock()
-        mock_registry.get_workflow.return_value = mock_workflow_func
+        # Create actual workflow definition
+        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions, workflow
+
+        class MockWorkflow:
+            @workflow.run
+            async def run(self):
+                return "test_result"
+
+        workflow_opts = WorkflowDefinitionOptions(name="test_workflow")
+        workflow_definition = WorkflowDefinition.wrap(MockWorkflow, workflow_opts)
+        mock_registry.get_workflow.return_value = workflow_definition
 
         # Create two different decision tasks
         task1 = Mock(spec=PollForDecisionTaskResponse)
@@ -333,10 +347,17 @@ async def test_respond_decision_task_completed_error(self, handler, sample_decis
     @pytest.mark.asyncio
     async def test_workflow_engine_creation_with_workflow_info(self, handler, sample_decision_task, mock_registry):
         """Test that WorkflowEngine is created with correct WorkflowInfo."""
-        mock_workflow_func = Mock()
-        mock_workflow_definition = Mock()
-        mock_workflow_definition.fn = mock_workflow_func
-        mock_registry.get_workflow.return_value = mock_workflow_definition
+        # Create actual workflow definition
+        from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions, workflow
+
+        class MockWorkflow:
+            @workflow.run
+            async def run(self):
+                return "test_result"
+
+        workflow_opts = WorkflowDefinitionOptions(name="test_workflow")
+        workflow_definition = WorkflowDefinition.wrap(MockWorkflow, workflow_opts)
+        mock_registry.get_workflow.return_value = workflow_definition
 
         mock_engine = Mock(spec=WorkflowEngine)
         mock_engine._is_workflow_complete = False  # Add missing attribute
@@ -363,4 +384,4 @@ async def test_workflow_engine_creation_with_workflow_info(self, handler, sample
                 call_args = mock_workflow_engine_class.call_args
                 assert call_args[1]['info'] is not None
                 assert call_args[1]['client'] == handler._client
-                assert call_args[1]['workflow_func'] == mock_workflow_func
+                assert call_args[1]['workflow_definition'] == workflow_definition

tests/cadence/worker/test_decision_task_handler_integration.py

@@ -13,6 +13,7 @@
 from cadence.api.v1.decision_pb2 import Decision
 from cadence.worker._decision_task_handler import DecisionTaskHandler
 from cadence.worker._registry import Registry
+from cadence.workflow import workflow
 from cadence.client import Client
 
 
@@ -37,9 +38,11 @@ def registry(self):
         reg = Registry()
 
         @reg.workflow(name="test_workflow")
-        def test_workflow(input_data):
-            """Simple test workflow that returns the input."""
-            return f"processed: {input_data}"
+        class TestWorkflow:
+            @workflow.run
+            async def run(self, input_data):
+                """Simple test workflow that returns the input."""
+                return f"processed: {input_data}"
 
         return reg
 

tests/cadence/worker/test_decision_worker_integration.py

@@ -11,6 +11,7 @@
 from cadence.api.v1.history_pb2 import History, HistoryEvent, WorkflowExecutionStartedEventAttributes
 from cadence.worker._decision import DecisionWorker
 from cadence.worker._registry import Registry
+from cadence.workflow import workflow
 from cadence.client import Client
 
 
@@ -34,12 +35,14 @@ def mock_client(self):
     def registry(self):
         """Create a registry with a test workflow."""
         reg = Registry()
-        
+
         @reg.workflow
-        def test_workflow(input_data):
-            """Simple test workflow that returns the input."""
-            return f"processed: {input_data}"
-        
+        class TestWorkflow:
+            @workflow.run
+            async def run(self, input_data):
+                """Simple test workflow that returns the input."""
+                return f"processed: {input_data}"
+
         return reg
 
     @pytest.fixture
@@ -236,8 +239,10 @@ async def test_decision_worker_with_different_workflow_types(self, decision_work
         """Test decision worker with different workflow types."""
         # Add another workflow to the registry
         @registry.workflow
-        def another_workflow(input_data):
-            return f"another-processed: {input_data}"
+        class AnotherWorkflow:
+            @workflow.run
+            async def run(self, input_data):
+                return f"another-processed: {input_data}"
 
         # Create decision tasks for different workflow types
         task1 = self.create_mock_decision_task(workflow_type="test_workflow")