From 621dbccb5a55e1ed57cb0bd89e092c4a5426c97e Mon Sep 17 00:00:00 2001
From: Xing Wang <xingwang1991@gmail.com>
Date: Thu, 18 Apr 2024 12:57:51 +0200
Subject: [PATCH] supports use both node name and the node in the wait list.
 (#45)

---
 aiida_worktree/node.py  |  6 ++++--
 tests/test_node_wait.py | 20 ++++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_node_wait.py

diff --git a/aiida_worktree/node.py b/aiida_worktree/node.py
index 4672f449..5aed3da6 100644
--- a/aiida_worktree/node.py
+++ b/aiida_worktree/node.py
@@ -20,7 +20,7 @@ def __init__(self, **kwargs):
         """
         super().__init__(**kwargs)
         self.to_ctx = None
-        self.wait = None
+        self.wait = []
         self.process = None
         self.pk = None
         self._widget = NodeGraphWidget(
@@ -30,7 +30,9 @@ def __init__(self, **kwargs):
     def to_dict(self):
         ndata = super().to_dict()
         ndata["to_ctx"] = [] if self.to_ctx is None else self.to_ctx
-        ndata["wait"] = [] if self.wait is None else self.wait
+        ndata["wait"] = [
+            node if isinstance(node, str) else node.name for node in self.wait
+        ]
         ndata["process"] = self.process.uuid if self.process else None
         ndata["metadata"]["pk"] = self.process.pk if self.process else None
 
diff --git a/tests/test_node_wait.py b/tests/test_node_wait.py
new file mode 100644
index 00000000..5689b029
--- /dev/null
+++ b/tests/test_node_wait.py
@@ -0,0 +1,20 @@
+import aiida
+
+aiida.load_profile()
+
+
+def test_node_wait(decorated_add):
+    """Run simple calcfunction."""
+    from aiida_worktree import WorkTree
+
+    wt = WorkTree(name="test_node_wait")
+    add1 = wt.nodes.new(decorated_add, "add1", x=1, y=1)
+    add1.to_ctx = [["result", "sum1"]]
+    add2 = wt.nodes.new(decorated_add, "add2", x=2, y=2)
+    add2.to_ctx = [["result", "sum2"]]
+    add3 = wt.nodes.new(decorated_add, "add3", x="{{sum1}}", y="{{sum2}}")
+    add3.wait = ["add1", add2]
+    wt.submit(wait=True)
+    add3.ctime < add1.ctime
+    add3.ctime < add2.ctime
+    assert add3.outputs[0].value == 6