IWF-464: Add WaitForStateExecutionCompletion API (#85)

* IWF-464: Implement wait_for_state_execution_completion

* IWF-464: Fix test

* IWF-464: Add waitForKey feature

* IWF-464: Fix test

* IWF-464: Remove sleep from tests

* IWF-464: Fix test

* IWF-464: Change test

* IWF-464: Lint

* IWF-464: Simplify test

* IWF-464: Add test complexity

Author: lwolczynski

iwf/client.py

@@ -22,6 +22,7 @@
     WorkflowState,
     get_state_id,
     get_state_id_by_class,
+    get_state_execution_id,
     should_skip_wait_until,
 )
 from iwf.workflow_state_options import _to_idl_state_options
@@ -86,6 +87,13 @@ def start_workflow(
 
             unreg_opts.workflow_config_override = options.workflow_config_override
 
+            unreg_opts.wait_for_completion_state_ids = (
+                options.wait_for_completion_state_ids
+            )
+            unreg_opts.wait_for_completion_state_execution_ids = (
+                options.wait_for_completion_state_execution_ids
+            )
+
             # TODO: set initial search attributes here
 
         starting_state_id = None
@@ -317,6 +325,47 @@ def set_workflow_search_attributes(
             workflow_class, workflow_id, run_id, search_attributes
         )
 
+    """A long poll API to wait for the completion of the state.
+    Note 1 The state_completion to wait for is needed to registered on starting workflow due to limitation in https://github.com/indeedeng/iwf/issues/349
+    Note 2 The max polling time is configured in client_options (default to 10s)
+
+    Args:
+        state_class the state class.
+        workflow_id the workflowId
+        state_execution_number the state execution number. E.g. if it's 2, it means the 2nd execution of the state
+    """
+
+    def wait_for_state_execution_completion_with_state_execution_id(
+        self,
+        state_class: type[WorkflowState],
+        workflow_id: str,
+        state_execution_number: int = 1,
+    ):
+        state_execution_id = get_state_execution_id(state_class, state_execution_number)
+
+        self._unregistered_client.wait_for_state_execution_completion_with_state_execution_id(
+            workflow_id, state_execution_id
+        )
+
+    """A long poll API to wait for the completion of the state.
+    Note 1 The state_completion to wait for is needed to registered on starting workflow due to limitation in https://github.com/indeedeng/iwf/issues/349
+    Note 2 The max polling time is configured in client_options (default to 10s)
+
+    Args:
+        state_class the state class.
+        workflow_id the workflowId
+        wait_for_key key provided by the client and to identity workflow
+    """
+
+    def wait_for_state_execution_completion_with_wait_for_key(
+        self, state_class: type[WorkflowState], workflow_id: str, wait_for_key: str
+    ):
+        state_id = get_state_id_by_class(state_class)
+
+        self._unregistered_client.wait_for_state_execution_completion_with_wait_for_key(
+            workflow_id, state_id, wait_for_key
+        )
+
     def _do_set_workflow_search_attributes(
         self,
         workflow_class: type[ObjectWorkflow],

iwf/client_options.py

@@ -9,6 +9,7 @@ class ClientOptions:
     worker_url: str
     object_encoder: ObjectEncoder
     api_timeout: int = 60
+    long_poll_api_max_wait_time_seconds: int = 10
 
     @classmethod
     def local_default(cls):

iwf/state_decision.py

@@ -54,9 +54,14 @@ def single_next_state(
         state: Union[str, type[WorkflowState]],
         state_input: Any = None,
         state_options_override: Optional[WorkflowStateOptions] = None,
+        wait_for_key: Optional[str] = None,
     ) -> StateDecision:
         return StateDecision(
-            [StateMovement.create(state, state_input, state_options_override)]
+            [
+                StateMovement.create(
+                    state, state_input, state_options_override, wait_for_key
+                )
+            ]
         )
 
     @classmethod

iwf/state_movement.py

@@ -36,6 +36,7 @@ class StateMovement:
     state_id: str
     state_input: Any = None
     state_options_override: Optional[WorkflowStateOptions] = None
+    wait_for_key: Optional[str] = None
 
     dead_end: typing.ClassVar[StateMovement]
 
@@ -57,6 +58,7 @@ def create(
         state: Union[str, type[WorkflowState]],
         state_input: Any = None,
         state_options_override: Optional[WorkflowStateOptions] = None,
+        wait_for_key: Optional[str] = None,
     ) -> StateMovement:
         if isinstance(state, str):
             state_id = state
@@ -68,7 +70,9 @@ 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, state_options_override)
+        return StateMovement(
+            state_id, state_input, state_options_override, wait_for_key
+        )
 
 
 StateMovement.dead_end = StateMovement(dead_end_sys_state_id)
@@ -99,4 +103,7 @@ def _to_idl_state_movement(
         )
 
         idl_movement.state_options = idl_state_options
+
+        if movement.wait_for_key is not None:
+            idl_movement.wait_for_key = movement.wait_for_key
     return idl_movement

iwf/tests/test_persistence_search_attributes.py

@@ -1,9 +1,10 @@
 import inspect
 import time
 import unittest
+from time import sleep
 
 from iwf.client import Client
-from iwf.command_request import CommandRequest, TimerCommand
+from iwf.command_request import CommandRequest
 from iwf.command_results import CommandResults
 from iwf.communication import Communication
 from iwf.iwf_api.models import SearchAttributeValueType
@@ -14,25 +15,24 @@
 from iwf.tests.worker_server import registry
 from iwf.workflow import ObjectWorkflow
 from iwf.workflow_context import WorkflowContext
+from iwf.workflow_options import WorkflowOptions
 from iwf.workflow_state import T, WorkflowState
 
 sa_keyword_key = "CustomKeywordField"
-sa_text_key = "CustomTextField"
 sa_double_key = "CustomDoubleField"
 sa_int_key = "CustomIntField"
 sa_bool_key = "CustomBoolField"
 sa_datetime_key = "CustomDatetimeField"
 sa_keyword_array_key = "CustomKeywordArrayField"
 
 sa_keyword: str = "keyword"
-sa_text: str = "text"
 sa_double: float = 2.34
 sa_int: int = 234
+sa_bool: bool = False
 sa_datetime: str = "2024-11-12T16:00:01.731455544-08:00"
 sa_keyword_array: list[str] = ["keyword-1", "keyword-2"]
 
 final_sa_keyword: str = "final_keyword"
-final_sa_text = None
 final_sa_int: int = 567
 final_sa_bool: bool = False
 final_sa_datetime: str = "2024-12-13T16:00:01.731455544-08:00"
@@ -47,14 +47,6 @@ def wait_until(
         persistence: Persistence,
         communication: Communication,
     ) -> CommandRequest:
-        persistence.set_search_attribute_keyword(sa_keyword_key, sa_keyword)
-        persistence.set_search_attribute_text(sa_text_key, sa_text)
-        persistence.set_search_attribute_double(sa_double_key, sa_double)
-        persistence.set_search_attribute_int64(sa_int_key, sa_int)
-        persistence.set_search_attribute_datetime(sa_datetime_key, sa_datetime)
-        persistence.set_search_attribute_keyword_array(
-            sa_keyword_array_key, sa_keyword_array
-        )
         return CommandRequest.empty()
 
     def execute(
@@ -65,6 +57,14 @@ def execute(
         persistence: Persistence,
         communication: Communication,
     ) -> StateDecision:
+        persistence.set_search_attribute_keyword(sa_keyword_key, sa_keyword)
+        persistence.set_search_attribute_double(sa_double_key, sa_double)
+        persistence.set_search_attribute_boolean(sa_bool_key, sa_bool)
+        persistence.set_search_attribute_keyword_array(
+            sa_keyword_array_key, sa_keyword_array
+        )
+        persistence.set_search_attribute_int64(sa_int_key, sa_int)
+        persistence.set_search_attribute_datetime(sa_datetime_key, sa_datetime)
         return StateDecision.single_next_state(SearchAttributeState2)
 
 
@@ -76,9 +76,7 @@ def wait_until(
         persistence: Persistence,
         communication: Communication,
     ) -> CommandRequest:
-        return CommandRequest.for_all_command_completed(
-            TimerCommand.by_seconds(7),
-        )
+        return CommandRequest.empty()
 
     def execute(
         self,
@@ -88,14 +86,15 @@ def execute(
         persistence: Persistence,
         communication: Communication,
     ) -> StateDecision:
+        # Delay updating search attributes to allow for the first assertion
+        sleep(1)
         persistence.set_search_attribute_keyword(sa_keyword_key, final_sa_keyword)
-        persistence.set_search_attribute_text(sa_text_key, final_sa_text)
-        persistence.set_search_attribute_int64(sa_int_key, final_sa_int)
         persistence.set_search_attribute_boolean(sa_bool_key, final_sa_bool)
-        persistence.set_search_attribute_datetime(sa_datetime_key, final_sa_datetime)
         persistence.set_search_attribute_keyword_array(
             sa_keyword_array_key, final_sa_keyword_array
         )
+        persistence.set_search_attribute_int64(sa_int_key, final_sa_int)
+        persistence.set_search_attribute_datetime(sa_datetime_key, final_sa_datetime)
         return StateDecision.graceful_complete_workflow()
 
 
@@ -110,23 +109,20 @@ def get_persistence_schema(self) -> PersistenceSchema:
             PersistenceField.search_attribute_def(
                 sa_keyword_key, SearchAttributeValueType.KEYWORD
             ),
-            PersistenceField.search_attribute_def(
-                sa_text_key, SearchAttributeValueType.TEXT
-            ),
             PersistenceField.search_attribute_def(
                 sa_double_key, SearchAttributeValueType.DOUBLE
             ),
             PersistenceField.search_attribute_def(
-                sa_int_key, SearchAttributeValueType.INT
+                sa_bool_key, SearchAttributeValueType.BOOL
             ),
             PersistenceField.search_attribute_def(
-                sa_bool_key, SearchAttributeValueType.BOOL
+                sa_keyword_array_key, SearchAttributeValueType.KEYWORD_ARRAY
             ),
             PersistenceField.search_attribute_def(
-                sa_datetime_key, SearchAttributeValueType.DATETIME
+                sa_int_key, SearchAttributeValueType.INT
             ),
             PersistenceField.search_attribute_def(
-                sa_keyword_array_key, SearchAttributeValueType.KEYWORD_ARRAY
+                sa_datetime_key, SearchAttributeValueType.DATETIME
             ),
         )
 
@@ -141,22 +137,27 @@ def setUpClass(cls):
     def test_persistence_search_attributes_workflow(self):
         wf_id = f"{inspect.currentframe().f_code.co_name}-{time.time_ns()}"
 
-        self.client.start_workflow(PersistenceSearchAttributesWorkflow, wf_id, 100)
+        wf_opts = WorkflowOptions()
+        wf_opts.add_wait_for_completion_state_ids(SearchAttributeState1)
+        self.client.start_workflow(
+            PersistenceSearchAttributesWorkflow, wf_id, 100, None, wf_opts
+        )
 
-        # TODO: Should be replaced with wait_for_state_execution_completed once implemented
-        # https://github.com/indeedeng/iwf-python-sdk/issues/48
-        time.sleep(5)
+        self.client.wait_for_state_execution_completion_with_state_execution_id(
+            SearchAttributeState1, wf_id
+        )
 
         returned_search_attributes = self.client.get_all_search_attributes(
-            PersistenceSearchAttributesWorkflow, wf_id
+            PersistenceSearchAttributesWorkflow,
+            wf_id,
         )
 
         expected_search_attributes = dict()
         expected_search_attributes[sa_keyword_key] = sa_keyword
-        expected_search_attributes[sa_text_key] = sa_text
         expected_search_attributes[sa_double_key] = sa_double
-        expected_search_attributes[sa_int_key] = sa_int
+        expected_search_attributes[sa_bool_key] = sa_bool
         expected_search_attributes[sa_keyword_array_key] = sa_keyword_array
+        expected_search_attributes[sa_int_key] = sa_int
         expected_search_attributes[sa_datetime_key] = (
             "2024-11-13T00:00:01.731455544Z"  # This is a bug. The iwf-server always returns utc time. See https://github.com/indeedeng/iwf/issues/261
             # "2024-11-12T18:00:01.731455544-06:00"
@@ -166,14 +167,8 @@ def test_persistence_search_attributes_workflow(self):
 
         self.client.wait_for_workflow_completion(wf_id)
 
-        final_returned_search_attributes = self.client.get_all_search_attributes(
-            PersistenceSearchAttributesWorkflow,
-            wf_id,
-        )
-
         final_expected_search_attributes = dict()
         final_expected_search_attributes[sa_keyword_key] = final_sa_keyword
-        final_expected_search_attributes[sa_text_key] = ""
         final_expected_search_attributes[sa_double_key] = sa_double
         final_expected_search_attributes[sa_int_key] = final_sa_int
         final_expected_search_attributes[sa_bool_key] = final_sa_bool
@@ -183,4 +178,9 @@ def test_persistence_search_attributes_workflow(self):
             # "2024-12-13T18:00:01.731455544-06:00"
         )
 
+        final_returned_search_attributes = self.client.get_all_search_attributes(
+            PersistenceSearchAttributesWorkflow,
+            wf_id,
+        )
+
         assert final_expected_search_attributes == final_returned_search_attributes