IWF-543: Support overriding state options (#79)

* IWF-543: Support overriding state options

* IWF-543: Fix types

* IWF-543: Add test

* IWF-543: Lint

* IWF-543: Lint

Author: lwolczynski

iwf/state_decision.py

@@ -3,13 +3,14 @@
 import typing
 
 from iwf.iwf_api.models import WorkflowConditionalClose, WorkflowConditionalCloseType
+from iwf.workflow_state_options import WorkflowStateOptions
 
 if typing.TYPE_CHECKING:
     from iwf.registry import Registry
     from iwf.workflow_state import WorkflowState
 
 from dataclasses import dataclass
-from typing import Any, List, Union
+from typing import Any, List, Union, Optional
 
 from iwf.iwf_api.models.state_decision import StateDecision as IdlStateDecision
 
@@ -49,9 +50,14 @@ def force_fail_workflow(cls, output: Any = None) -> StateDecision:
 
     @classmethod
     def single_next_state(
-        cls, state: Union[str, type[WorkflowState]], state_input: Any = None
+        cls,
+        state: Union[str, type[WorkflowState]],
+        state_input: Any = None,
+        state_options_override: Optional[WorkflowStateOptions] = None,
     ) -> StateDecision:
-        return StateDecision([StateMovement.create(state, state_input)])
+        return StateDecision(
+            [StateMovement.create(state, state_input, state_options_override)]
+        )
 
     @classmethod
     def multi_next_states(

iwf/state_movement.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import typing
-from typing import Union
+from typing import Union, Optional
 
 from iwf.errors import WorkflowDefinitionError
 
@@ -14,12 +14,12 @@
 from dataclasses import dataclass
 from typing import Any
 
+
 from iwf.iwf_api.models.state_movement import StateMovement as IdlStateMovement
 
 from iwf.object_encoder import ObjectEncoder
 
-from iwf.workflow_state_options import _to_idl_state_options
-
+from iwf.workflow_state_options import _to_idl_state_options, WorkflowStateOptions
 
 reserved_state_id_prefix = "_SYS_"
 
@@ -35,6 +35,7 @@
 class StateMovement:
     state_id: str
     state_input: Any = None
+    state_options_override: Optional[WorkflowStateOptions] = None
 
     dead_end: typing.ClassVar[StateMovement]
 
@@ -52,7 +53,10 @@ def force_fail_workflow(cls, output: Any = None) -> StateMovement:
 
     @classmethod
     def create(
-        cls, state: Union[str, type[WorkflowState]], state_input: Any = None
+        cls,
+        state: Union[str, type[WorkflowState]],
+        state_input: Any = None,
+        state_options_override: Optional[WorkflowStateOptions] = None,
     ) -> StateMovement:
         if isinstance(state, str):
             state_id = state
@@ -64,7 +68,7 @@ def create(
             state_id = get_state_id_by_class(state)
         if state_id.startswith(reserved_state_id_prefix):
             raise WorkflowDefinitionError("cannot use reserved stateId")
-        return StateMovement(state_id, state_input)
+        return StateMovement(state_id, state_input, state_options_override)
 
 
 StateMovement.dead_end = StateMovement(dead_end_sys_state_id)
@@ -83,9 +87,14 @@ def _to_idl_state_movement(
             should_skip_wait_until,
         )
 
+        if movement.state_options_override is not None:
+            options = movement.state_options_override
+        else:
+            options = state.get_state_options()
+
         idl_state_options = _to_idl_state_options(
             should_skip_wait_until(state),
-            state.get_state_options(),
+            options,
             registry.get_state_store(wf_type),
         )
 

iwf/tests/test_workflow_state_options_override.py

@@ -0,0 +1,117 @@
+import inspect
+import time
+import unittest
+
+from iwf.client import Client
+from iwf.command_request import CommandRequest
+from iwf.command_results import CommandResults
+from iwf.communication import Communication
+from iwf.persistence_schema import PersistenceSchema, PersistenceField
+from iwf.workflow_options import WorkflowOptions
+
+from iwf.iwf_api.models import RetryPolicy, WaitUntilApiFailurePolicy, IDReusePolicy
+from iwf.persistence import Persistence
+from iwf.state_decision import StateDecision
+from iwf.state_schema import StateSchema
+from iwf.tests.worker_server import registry
+from iwf.workflow import ObjectWorkflow
+from iwf.workflow_context import WorkflowContext
+from iwf.workflow_state import T, WorkflowState
+from iwf.workflow_state_options import WorkflowStateOptions
+
+output_da = "output_da"
+
+
+class InitState(WorkflowState[str]):
+    def wait_until(
+        self,
+        ctx: WorkflowContext,
+        input: T,
+        persistence: Persistence,
+        communication: Communication,
+    ) -> CommandRequest:
+        persistence.set_data_attribute(
+            output_da, str(input) + "_InitState_waitUntil_completed"
+        )
+        return CommandRequest.empty()
+
+    def execute(
+        self,
+        ctx: WorkflowContext,
+        input: T,
+        command_results: CommandResults,
+        persistence: Persistence,
+        communication: Communication,
+    ) -> StateDecision:
+        data = persistence.get_data_attribute(output_da)
+        data += "_InitState_execute_completed"
+        return StateDecision.single_next_state(
+            NonInitState,
+            data,
+            WorkflowStateOptions(
+                wait_until_api_retry_policy=RetryPolicy(maximum_attempts=2),
+                proceed_to_execute_when_wait_until_retry_exhausted=WaitUntilApiFailurePolicy.PROCEED_ON_FAILURE,
+            ),
+        )
+
+
+class NonInitState(WorkflowState[str]):
+    def wait_until(
+        self,
+        ctx: WorkflowContext,
+        input: T,
+        persistence: Persistence,
+        communication: Communication,
+    ) -> CommandRequest:
+        raise RuntimeError("test failure")
+
+    def execute(
+        self,
+        ctx: WorkflowContext,
+        input: T,
+        command_results: CommandResults,
+        persistence: Persistence,
+        communication: Communication,
+    ) -> StateDecision:
+        data = str(input) + "_NonInitState_execute_completed"
+        return StateDecision.graceful_complete_workflow(data)
+
+    def get_state_options(self) -> WorkflowStateOptions:
+        return WorkflowStateOptions(
+            wait_until_api_retry_policy=RetryPolicy(maximum_attempts=1),
+            proceed_to_execute_when_wait_until_retry_exhausted=WaitUntilApiFailurePolicy.FAIL_WORKFLOW_ON_FAILURE,
+        )
+
+
+class StateOptionsOverrideWorkflow(ObjectWorkflow):
+    def get_persistence_schema(self) -> PersistenceSchema:
+        return PersistenceSchema.create(
+            PersistenceField.data_attribute_def(output_da, str),
+        )
+
+    def get_workflow_states(self) -> StateSchema:
+        return StateSchema.with_starting_state(InitState(), NonInitState())
+
+
+class TestStateOptionsOverrideWorkflow(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        wf = StateOptionsOverrideWorkflow()
+        registry.add_workflow(wf)
+        cls.client = Client(registry)
+
+    def test_override(self):
+        wf_id = f"{inspect.currentframe().f_code.co_name}-{time.time_ns()}"
+        self.client.start_workflow(
+            StateOptionsOverrideWorkflow,
+            wf_id,
+            10,
+            "input",
+            WorkflowOptions(workflow_id_reuse_policy=IDReusePolicy.DISALLOW_REUSE),
+        )
+        output = self.client.wait_for_workflow_completion(wf_id)
+
+        assert (
+            output
+            == "input_InitState_waitUntil_completed_InitState_execute_completed_NonInitState_execute_completed"
+        )