Test dynamically registering workflows, steps, transactions in class (#440)

Test that we can register methods from within the __init__ function of a
non-DBOS class. This allows us to give dynamic names to the
workflow/step/transaction methods based on a given name if needed.

Author: qianl15

tests/classdefs.py

@@ -92,3 +92,36 @@ def test_recv_workflow(topic: str) -> str:
         msg3 = DBOS.recv(timeout_seconds=10)
         DBOSSendRecv.recv_counter += 1
         return "-".join([str(msg1), str(msg2), str(msg3)])
+
+
+class DBOSTestWrapperMethods:
+    # Test that we can register methods from within the __init__ function of a non-DBOS class
+    # This allows us to give dynamic names to the methods if needed
+
+    def __init__(self, name: str) -> None:
+        self.wf_counter: int = 0
+        self.step_counter: int = 0
+        self.txn_counter: int = 0
+
+        self.name = name
+        assert self.name is not None
+
+        @DBOS.step(name=f"{self.name}_test_step")
+        def wrapped_test_step(var: str) -> str:
+            self.step_counter += 1
+            return var
+
+        @DBOS.transaction(name=f"{self.name}_test_transaction")
+        def wrapped_test_transaction(var2: str) -> str:
+            rows = DBOS.sql_session.execute(sa.text("SELECT 1")).fetchall()
+            self.txn_counter += 1
+            return var2 + str(rows[0][0])
+
+        @DBOS.workflow(name=f"{self.name}_test_workflow")
+        def wrapped_test_workflow(var: str, var2: str) -> str:
+            self.wf_counter += 1
+            res = wrapped_test_transaction(var2)
+            res2 = wrapped_test_step(var)
+            return res + res2
+
+        self.test_workflow = wrapped_test_workflow

tests/test_singleton.py

@@ -21,7 +21,12 @@ def test_dbos_singleton(cleanup_test_databases: None) -> None:
 
     # Simulate an app that does some imports of its own code, then defines DBOS,
     #    then imports more
-    from tests.classdefs import DBOSSendRecv, DBOSTestClass, DBOSTestRoles
+    from tests.classdefs import (
+        DBOSSendRecv,
+        DBOSTestClass,
+        DBOSTestRoles,
+        DBOSTestWrapperMethods,
+    )
 
     dbos: DBOS = DBOS(config=default_config())
 
@@ -53,6 +58,26 @@ def test_dbos_singleton(cleanup_test_databases: None) -> None:
     wfhr: WorkflowHandle[str] = DBOS.retrieve_workflow(wh.get_workflow_id())
     assert wfhr.workflow_id == wh.get_workflow_id()
 
+    # Test wrapper class methods
+    wrapper_inst = DBOSTestWrapperMethods("dynamicconfig")
+    res = wrapper_inst.test_workflow("x", "y")
+    assert res == "y1x"
+
+    wh = DBOS.start_workflow(wrapper_inst.test_workflow, "x", "y")
+    assert wh.get_result() == "y1x"
+    stati = DBOS.get_workflow_status(wh.get_workflow_id())
+    assert stati
+    assert stati.config_name is None
+    assert stati.class_name is None
+    assert stati.name == "dynamicconfig_test_workflow"
+    wfhr = DBOS.retrieve_workflow(wh.get_workflow_id())
+    assert wfhr.workflow_id == wh.get_workflow_id()
+
+    stepi = DBOS.list_workflow_steps(wh.get_workflow_id())
+    assert len(stepi) == 2
+    assert stepi[0]["function_name"] == "dynamicconfig_test_transaction"
+    assert stepi[1]["function_name"] == "dynamicconfig_test_step"
+
     # Roles
 
     with DBOSContextEnsure():