Add test for timeout

tests/conftest.py

@@ -242,12 +242,49 @@ def wg_engine(decorated_add, add_code) -> WorkGraph:
 
 
 @pytest.fixture
-def finished_process_node():
-    """Return a finished process node."""
-
-    node = WorkflowNode()
-    node.set_process_state("finished")
-    node.set_exit_status(0)
-    node.seal()
-    node.store()
-    return node
+def create_process_node():
+    """Return a process node."""
+
+    def process_node(state="finished", exit_status=0):
+        """Return a finished process node."""
+
+        node = WorkflowNode()
+        node.set_process_state(state)
+        node.set_exit_status(exit_status)
+        node.seal()
+        node.store()
+        return node
+
+    return process_node
+
+
+@pytest.fixture
+def create_workgraph_process_node():
+    """Return a process node."""
+
+    def process_node(state="finished", exit_status=0):
+        """Return a finished process node."""
+        from aiida_workgraph.engine.workgraph import WorkGraphEngine
+
+        process = WorkGraphEngine(
+            inputs={
+                "wg": {
+                    "name": "test",
+                    "uuid": "",
+                    "tasks": {},
+                    "links": [],
+                    "ctrl_links": [],
+                    "state": "",
+                    "error_handlers": {},
+                    "context": {},
+                }
+            }
+        )
+        node = process.node
+        node.set_process_state(state)
+        node.set_exit_status(exit_status)
+        node.seal()
+        node.store()
+        return node
+
+    return process_node

tests/test_action.py

@@ -39,10 +39,10 @@ def test_pause_play_task(wg_calcjob):
     assert wg.tasks.add2.outputs.sum.value == 9
 
 
-def test_pause_play_error_handler(wg_calcjob, finished_process_node):
+def test_pause_play_error_handler(wg_calcjob, create_process_node):
     wg = wg_calcjob
     wg.name = "test_pause_play_error_handler"
-    wg.process = finished_process_node
+    wg.process = create_process_node(state="finished", exit_status=0)
     try:
         wg.pause_tasks(["add1"])
     except Exception as e:

tests/test_workgraph.py

@@ -181,3 +181,14 @@ def test_workgraph_group_outputs(decorated_add):
     wg.run()
     assert wg.process.outputs.sum.value == 5
     # assert wg.process.outputs.add1.result.value == 5
+
+
+@pytest.mark.usefixtures("started_daemon_client")
+def test_wait_timeout(create_workgraph_process_node):
+    wg = WorkGraph()
+    wg.process = create_workgraph_process_node(state="running")
+    with pytest.raises(
+        TimeoutError,
+        match="Timeout reached after 1 seconds while waiting for the WorkGraph:",
+    ):
+        wg.wait(timeout=1, interval=1)