IWF-686: Standardize unit test implementation (#102)

* IWF-686: Standardize unit test implementation

* IWF-686: Refactor

Author: lwolczynski

iwf/tests/__init__.py

@@ -0,0 +1,56 @@
+from iwf.registry import Registry
+from iwf.tests.workflows.abnormal_exit_workflow import AbnormalExitWorkflow
+from iwf.tests.workflows.basic_workflow import BasicWorkflow
+from iwf.tests.workflows.conditional_complete_workflow import (
+    ConditionalCompleteWorkflow,
+)
+from iwf.tests.workflows.describe_workflow import DescribeWorkflow
+from iwf.tests.workflows.internal_channel_workflow import InternalChannelWorkflow
+from iwf.tests.workflows.internal_channel_workflow_with_no_prefix_channel import (
+    InternalChannelWorkflowWithNoPrefixChannel,
+)
+from iwf.tests.workflows.persistence_data_attributes_workflow import (
+    PersistenceDataAttributesWorkflow,
+)
+from iwf.tests.workflows.persistence_search_attributes_workflow import (
+    PersistenceSearchAttributesWorkflow,
+)
+from iwf.tests.workflows.persistence_state_execution_local_workflow import (
+    PersistenceStateExecutionLocalWorkflow,
+)
+from iwf.tests.workflows.recovery_workflow import RecoveryWorkflow
+from iwf.tests.workflows.rpc_workflow import RPCWorkflow
+from iwf.tests.workflows.state_options_override_workflow import (
+    StateOptionsOverrideWorkflow,
+)
+from iwf.tests.workflows.timer_workflow import TimerWorkflow
+from iwf.tests.workflows.wait_for_state_with_state_execution_id_workflow import (
+    WaitForStateWithStateExecutionIdWorkflow,
+)
+from iwf.tests.workflows.wait_for_state_with_wait_for_key_workflow import (
+    WaitForStateWithWaitForKeyWorkflow,
+)
+from iwf.tests.workflows.wait_internal_channel_workflow import (
+    WaitInternalChannelWorkflow,
+)
+from iwf.tests.workflows.wait_signal_workflow import WaitSignalWorkflow
+
+registry = Registry()
+
+registry.add_workflow(AbnormalExitWorkflow())
+registry.add_workflow(BasicWorkflow())
+registry.add_workflow(ConditionalCompleteWorkflow())
+registry.add_workflow(DescribeWorkflow())
+registry.add_workflow(InternalChannelWorkflow())
+registry.add_workflow(InternalChannelWorkflowWithNoPrefixChannel())
+registry.add_workflow(PersistenceDataAttributesWorkflow())
+registry.add_workflow(PersistenceSearchAttributesWorkflow())
+registry.add_workflow(PersistenceStateExecutionLocalWorkflow())
+registry.add_workflow(RecoveryWorkflow())
+registry.add_workflow(RPCWorkflow())
+registry.add_workflow(TimerWorkflow())
+registry.add_workflow(StateOptionsOverrideWorkflow())
+registry.add_workflow(WaitForStateWithStateExecutionIdWorkflow())
+registry.add_workflow(WaitForStateWithWaitForKeyWorkflow())
+registry.add_workflow(WaitInternalChannelWorkflow())
+registry.add_workflow(WaitSignalWorkflow())

iwf/tests/test_abnormal_exit_workflow.py

@@ -1,65 +1,34 @@
 import inspect
 import time
 import unittest
-from typing import Union
 
 from iwf.client import Client
-from iwf.command_results import CommandResults
-from iwf.communication import Communication
 from iwf.errors import WorkflowFailed
-from iwf.iwf_api.models import RetryPolicy
 from iwf.iwf_api.models.id_reuse_policy import IDReusePolicy
-from iwf.persistence import Persistence
-from iwf.state_decision import StateDecision
-from iwf.state_schema import StateSchema
-from iwf.tests.test_basic_workflow import BasicWorkflow
 from iwf.tests.worker_server import registry
-from iwf.workflow import ObjectWorkflow
-from iwf.workflow_context import WorkflowContext
+from iwf.tests.workflows.abnormal_exit_workflow import AbnormalExitWorkflow
+from iwf.tests.workflows.basic_workflow import BasicWorkflow
 from iwf.workflow_options import WorkflowOptions
-from iwf.workflow_state import T, WorkflowState
-from iwf.workflow_state_options import WorkflowStateOptions
-
-
-class AbnormalExitState1(WorkflowState[Union[int, str]]):
-    def execute(
-        self,
-        ctx: WorkflowContext,
-        input: T,
-        command_results: CommandResults,
-        persistence: Persistence,
-        communication: Communication,
-    ) -> StateDecision:
-        raise RuntimeError("abnormal exit state")
-
-    def get_state_options(self) -> WorkflowStateOptions:
-        return WorkflowStateOptions(
-            execute_api_retry_policy=RetryPolicy(maximum_attempts=1)
-        )
-
-
-class AbnormalExitWorkflow(ObjectWorkflow):
-    def get_workflow_states(self) -> StateSchema:
-        return StateSchema.with_starting_state(AbnormalExitState1())
-
-
-abnormal_exit_wf = AbnormalExitWorkflow()
-registry.add_workflow(abnormal_exit_wf)
-client = Client(registry)
 
 
 class TestAbnormalWorkflow(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.client = Client(registry)
+
     def test_abnormal_exit_workflow(self):
         wf_id = f"{inspect.currentframe().f_code.co_name}-{time.time_ns()}"
-        startOptions = WorkflowOptions(
+        start_options = WorkflowOptions(
             workflow_id_reuse_policy=IDReusePolicy.ALLOW_IF_PREVIOUS_EXITS_ABNORMALLY
         )
 
-        client.start_workflow(AbnormalExitWorkflow, wf_id, 100, "input", startOptions)
+        self.client.start_workflow(
+            AbnormalExitWorkflow, wf_id, 100, "input", start_options
+        )
         with self.assertRaises(WorkflowFailed):
-            client.get_simple_workflow_result_with_wait(wf_id, str)
+            self.client.wait_for_workflow_completion(wf_id, str)
 
         # Starting a workflow with the same ID should be allowed since the previous failed abnormally
-        client.start_workflow(BasicWorkflow, wf_id, 100, "input", startOptions)
-        res = client.get_simple_workflow_result_with_wait(wf_id, str)
+        self.client.start_workflow(BasicWorkflow, wf_id, 100, "input", start_options)
+        res = self.client.wait_for_workflow_completion(wf_id, str)
         assert res == "done"

iwf/tests/test_basic_workflow.py

@@ -1,74 +1,20 @@
 import inspect
 import time
 import unittest
-from typing import Union
 
 from iwf.client import Client
-from iwf.command_request import CommandRequest, TimerCommand
-from iwf.command_results import CommandResults
-from iwf.communication import Communication
 from iwf.errors import WorkflowAlreadyStartedError
 from iwf.iwf_api.models import WorkflowAlreadyStartedOptions
-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.tests.workflows.basic_workflow import BasicWorkflow
 from iwf.workflow_options import WorkflowOptions
-from iwf.workflow_state import T, WorkflowState
-
-
-class State1(WorkflowState[Union[int, str]]):
-    def wait_until(
-        self,
-        ctx: WorkflowContext,
-        input: T,
-        persistence: Persistence,
-        communication: Communication,
-    ) -> CommandRequest:
-        if input != "input":
-            raise RuntimeError("input is incorrect")
-        return CommandRequest.for_all_command_completed(
-            TimerCommand.by_seconds(1),
-        )
-
-    def execute(
-        self,
-        ctx: WorkflowContext,
-        input: T,
-        command_results: CommandResults,
-        persistence: Persistence,
-        communication: Communication,
-    ) -> StateDecision:
-        if input != "input":
-            raise RuntimeError("input is incorrect")
-        return StateDecision.single_next_state(State2)
-
-
-class State2(WorkflowState[None]):
-    def execute(
-        self,
-        ctx: WorkflowContext,
-        input: T,
-        command_results: CommandResults,
-        persistence: Persistence,
-        communication: Communication,
-    ) -> StateDecision:
-        return StateDecision.graceful_complete_workflow("done")
-
-
-class BasicWorkflow(ObjectWorkflow):
-    def get_workflow_states(self) -> StateSchema:
-        return StateSchema.with_starting_state(State1(), State2())
-
-
-hello_wf = BasicWorkflow()
-registry.add_workflow(hello_wf)
-client = Client(registry)
 
 
 class TestWorkflowErrors(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.client = Client(registry)
+
     def test_basic_workflow(self):
         original_request_id = "1"
         later_request_id = "2"
@@ -85,12 +31,12 @@ def test_basic_workflow(self):
             workflow_already_started_options_1
         )
 
-        wf_run_id = client.start_workflow(
+        wf_run_id = self.client.start_workflow(
             BasicWorkflow, wf_id, 100, "input", start_options_1
         )
         assert wf_run_id
 
-        wf_run_id = client.start_workflow(
+        wf_run_id = self.client.start_workflow(
             BasicWorkflow, wf_id, 100, "input", start_options_1
         )
         assert wf_run_id
@@ -106,10 +52,10 @@ def test_basic_workflow(self):
         )
 
         with self.assertRaises(WorkflowAlreadyStartedError):
-            wf_run_id = client.start_workflow(
+            wf_run_id = self.client.start_workflow(
                 BasicWorkflow, wf_id, 100, "input", start_options_2
             )
             assert wf_run_id
 
-        res = client.wait_for_workflow_completion(wf_id, str)
+        res = self.client.wait_for_workflow_completion(wf_id, str)
         assert res == "done"

iwf/tests/test_conditional_complete.py

@@ -3,123 +3,40 @@
 import unittest
 
 from iwf.client import Client
-from iwf.command_request import (
-    CommandRequest,
-    InternalChannelCommand,
-    SignalChannelCommand,
-)
-from iwf.command_results import CommandResults
-from iwf.communication import Communication
-from iwf.communication_schema import CommunicationMethod, CommunicationSchema
-from iwf.persistence import Persistence
-from iwf.persistence_schema import PersistenceField, PersistenceSchema
-from iwf.rpc import rpc
-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
-
-test_signal_channel = "test-1"
-test_internal_channel = "test-2"
-
-da_counter = "counter"
-
-
-class WaitState(WorkflowState[bool]):
-    def wait_until(
-        self,
-        ctx: WorkflowContext,
-        use_signal: T,
-        persistence: Persistence,
-        communication: Communication,
-    ) -> CommandRequest:
-        if use_signal:
-            return CommandRequest.for_all_command_completed(
-                SignalChannelCommand.by_name(test_signal_channel),
-            )
-        else:
-            return CommandRequest.for_all_command_completed(
-                InternalChannelCommand.by_name(test_internal_channel),
-            )
-
-    def execute(
-        self,
-        ctx: WorkflowContext,
-        use_signal: T,
-        command_results: CommandResults,
-        persistence: Persistence,
-        communication: Communication,
-    ) -> StateDecision:
-        counter = persistence.get_data_attribute(da_counter)
-        if counter is None:
-            counter = 0
-        counter += 1
-        persistence.set_data_attribute(da_counter, counter)
-
-        if ctx.state_execution_id == "WaitState-1":
-            # wait for 3 seconds so that the channel can have a new message
-            time.sleep(3)
-        if use_signal:
-            return StateDecision.force_complete_if_signal_channel_empty_or_else(
-                test_signal_channel, counter, WaitState, use_signal
-            )
-        else:
-            return StateDecision.force_complete_if_internal_channel_empty_or_else(
-                test_internal_channel, counter, WaitState, use_signal
-            )
-
-
-class ConditionalCompleteWorkflow(ObjectWorkflow):
-    def get_communication_schema(self) -> CommunicationSchema:
-        return CommunicationSchema.create(
-            CommunicationMethod.signal_channel_def(test_signal_channel, int),
-            CommunicationMethod.internal_channel_def(test_internal_channel, int),
-        )
-
-    def get_persistence_schema(self) -> PersistenceSchema:
-        return PersistenceSchema.create(
-            PersistenceField.data_attribute_def(da_counter, int),
-        )
-
-    def get_workflow_states(self) -> StateSchema:
-        return StateSchema.with_starting_state(WaitState())
-
-    @rpc()
-    def test_rpc_publish_channel(self, com: Communication):
-        com.publish_to_internal_channel(test_internal_channel, 0)
+from iwf.tests.workflows.conditional_complete_workflow import (
+    ConditionalCompleteWorkflow,
+    test_signal_channel,
+)
 
 
 class TestConditionalComplete(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        wf = ConditionalCompleteWorkflow()
-        registry.add_workflow(wf)
+        cls.client = Client(registry)
 
     def test_internal_channel_conditional_complete(self):
         wf_id = f"{inspect.currentframe().f_code.co_name}-{time.time_ns()}"
-        self.do_test_conditional_workflow(wf_id, False)
+        do_test_conditional_workflow(self.client, wf_id, False)
 
     def test_signal_channel_conditional_complete(self):
         wf_id = f"{inspect.currentframe().f_code.co_name}-{time.time_ns()}"
-        self.do_test_conditional_workflow(wf_id, True)
+        do_test_conditional_workflow(self.client, wf_id, True)
 
-    def do_test_conditional_workflow(self, wf_id: str, use_signal: bool):
-        self.client = Client(registry)
 
-        self.client.start_workflow(ConditionalCompleteWorkflow, wf_id, 10, use_signal)
+def do_test_conditional_workflow(client: Client, wf_id: str, use_signal: bool):
+    client.start_workflow(ConditionalCompleteWorkflow, wf_id, 10, use_signal)
 
-        for x in range(3):
-            if use_signal:
-                self.client.signal_workflow(wf_id, test_signal_channel, 123)
-            else:
-                self.client.invoke_rpc(
-                    wf_id, ConditionalCompleteWorkflow.test_rpc_publish_channel
-                )
-            if x == 0:
-                # wait for a second so that the workflow is in execute state
-                time.sleep(1)
+    for x in range(3):
+        if use_signal:
+            client.signal_workflow(wf_id, test_signal_channel, 123)
+        else:
+            client.invoke_rpc(
+                wf_id, ConditionalCompleteWorkflow.test_rpc_publish_channel
+            )
+        if x == 0:
+            # wait for a second so that the workflow is in execute state
+            time.sleep(1)
 
-        res = self.client.get_simple_workflow_result_with_wait(wf_id)
-        assert res == 3
+    res = client.wait_for_workflow_completion(wf_id)
+    assert res == 3