IWF-356: Add test

iwf/tests/test_basic_workflow.py

@@ -6,14 +6,12 @@
 from iwf.command_request import CommandRequest
 from iwf.command_results import CommandResults
 from iwf.communication import Communication
-from iwf.iwf_api.models import WorkflowConfig
 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_options import WorkflowOptions
 from iwf.workflow_state import T, WorkflowState
 
 
@@ -67,13 +65,11 @@ def get_workflow_states(self) -> StateSchema:
 def test_basic_workflow():
     wf_id = f"{inspect.currentframe().f_code.co_name}-{time.time_ns()}"
 
-    # TODO: Add an override and verify with WebUI
     client.start_workflow(
         BasicWorkflow,
         wf_id,
         100,
         "input",
-        WorkflowOptions(workflow_config_override=WorkflowConfig()),
     )
     res = client.get_simple_workflow_result_with_wait(wf_id, str)
     assert res == "done"

iwf/tests/test_persistence.py

@@ -1,5 +1,6 @@
 import inspect
 import time
+import unittest
 
 from iwf.client import Client
 from iwf.command_request import CommandRequest
@@ -12,12 +13,22 @@
 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
+from iwf.rpc import rpc
+
+initial_da_1 = "initial_da_1"
+initial_da_value_1 = "value_1"
+initial_da_2 = "initial_da_2"
+initial_da_value_2 = "value_2"
 
 test_da_1 = "test_da_1"
 test_da_2 = "test_da_2"
 
-# TODO: Add a test for initial data attributes
+final_test_da_value_1 = "1234"
+final_test_da_value_2 = 1234
+final_initial_da_value_1 = initial_da_value_1
+final_initial_da_value_2 = "no-more-init"
 
 
 class DataAttributeRWState(WorkflowState[None]):
@@ -46,8 +57,9 @@ def execute(
         assert da1 == "123"
         assert da2 == 123
 
-        persistence.set_data_attribute(test_da_1, "1234")
-        persistence.set_data_attribute(test_da_2, 1234)
+        persistence.set_data_attribute(test_da_1, final_test_da_value_1)
+        persistence.set_data_attribute(test_da_2, final_test_da_value_2)
+        persistence.set_data_attribute(initial_da_2, final_initial_da_value_2)
         return StateDecision.graceful_complete_workflow()
 
 
@@ -57,19 +69,46 @@ def get_workflow_states(self) -> StateSchema:
 
     def get_persistence_schema(self) -> PersistenceSchema:
         return PersistenceSchema.create(
+            PersistenceField.data_attribute_def(initial_da_1, str),
+            PersistenceField.data_attribute_def(initial_da_2, str),
             PersistenceField.data_attribute_def(test_da_1, str),
             PersistenceField.data_attribute_def(test_da_2, int),
         )
 
+    @rpc()
+    def test_persistence_read(self, pers: Persistence):
+        return (
+            pers.get_data_attribute(initial_da_1),
+            pers.get_data_attribute(initial_da_2),
+            pers.get_data_attribute(test_da_1),
+            pers.get_data_attribute(test_da_2),
+        )
+
 
-wf = PersistenceWorkflow()
-registry.add_workflow(wf)
-client = Client(registry)
+class TestPersistence(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        wf = PersistenceWorkflow()
+        registry.add_workflow(wf)
+        cls.client = Client(registry)
 
+    def test_persistence_workflow(self):
+        wf_id = f"{inspect.currentframe().f_code.co_name}-{time.time_ns()}"
+
+        start_options = WorkflowOptions(
+            initial_data_attributes={
+                initial_da_1: initial_da_value_1,
+                initial_da_2: initial_da_value_2,
+            },
+        )
 
-def test_persistence_workflow():
-    wf_id = f"{inspect.currentframe().f_code.co_name}-{time.time_ns()}"
+        self.client.start_workflow(PersistenceWorkflow, wf_id, 100, None, start_options)
+        self.client.get_simple_workflow_result_with_wait(wf_id, None)
 
-    client.start_workflow(PersistenceWorkflow, wf_id, 100, None)
-    client.get_simple_workflow_result_with_wait(wf_id, None)
-    # TODO use RPC to get the final result of the persistence updated in execute API
+        res = self.client.invoke_rpc(wf_id, PersistenceWorkflow.test_persistence_read)
+        assert res == [
+            final_initial_da_value_1,
+            final_initial_da_value_2,
+            final_test_da_value_1,
+            final_test_da_value_2,
+        ]