From 18397b9bfdbcf3544607e5f3ae2a161ff21eba71 Mon Sep 17 00:00:00 2001
From: superstar54 <xingwang1991@gmail.com>
Date: Fri, 13 Dec 2024 07:28:06 +0100
Subject: [PATCH] Add test for timeout

---
 tests/conftest.py       | 55 ++++++++++++++++++++++++++++++++++-------
 tests/test_action.py    |  4 +--
 tests/test_workgraph.py | 11 +++++++++
 3 files changed, 59 insertions(+), 11 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index e553dc25..fce5d217 100644
--- a/tests/conftest.py
+++ b/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
diff --git a/tests/test_action.py b/tests/test_action.py
index bca6dcdd..ba403877 100644
--- a/tests/test_action.py
+++ b/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:
diff --git a/tests/test_workgraph.py b/tests/test_workgraph.py
index 1badd386..4d6d6215 100644
--- a/tests/test_workgraph.py
+++ b/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)