Improve process_tree - node_view behavior

edan-bainglass · Nov 14, 2024 · 00b015f · 00b015f
1 parent f16abc2
commit 00b015f
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 32 deletions.
diff --git a/src/aiidalab_qe/app/result/__init__.py b/src/aiidalab_qe/app/result/__init__.py
@@ -8,9 +8,10 @@
     ProcessNodesTreeWidget,
     WizardAppWidgetStep,
 )
+from aiidalab_widgets_base.viewers import viewer as node_viewer
 
 from .model import ResultsStepModel
-from .viewer import WorkChainViewer, WorkChainViewerModel  # noqa: F401
+from .viewer import WorkChainViewer, WorkChainViewerModel
 
 PROCESS_COMPLETED = "<h4>Workflow completed successfully!</h4>"
 PROCESS_EXCEPTED = "<h4>Workflow is excepted!</h4>"
@@ -173,38 +174,40 @@ def _update_node_view(self, nodes, refresh=False):
             If True, the viewer will be refreshed.
             Occurs when user presses the "Update results" button.
         """
-        from aiidalab_widgets_base.viewers import viewer
 
         if not nodes:
             return
         # only show the first selected node
         node = nodes[0]
+
         # check if the viewer is already added
         if node.uuid in self.node_views and not refresh:
             self.node_view = self.node_views[node.uuid]
-        else:
-            kwargs = {}
-            if isinstance(node, orm.WorkChainNode):
-                model = WorkChainViewerModel()
-                ipw.dlink(
-                    (self._model, "process_uuid"),
-                    (model, "process_uuid"),
-                )
-                ipw.dlink(
-                    (self._model, "monitor_counter"),
-                    (model, "monitor_counter"),
-                )
-                kwargs["model"] = model
-            self.node_view = viewer(node, **kwargs)
-            # TODO this should work - fix!
-            # if isinstance(self.node_view, WorkChainViewer):
-            #     self.node_view.render()
+        elif not isinstance(node, orm.WorkChainNode):
+            self.node_view = node_viewer(node)
             self.node_views[node.uuid] = self.node_view
+        elif node.process_label == "QeAppWorkChain":
+            self._create_workchain_viewer(node)
+
         self.process_status.children = [
             self.process_tree,
             self.node_view,
         ]
 
+    def _create_workchain_viewer(self, node):
+        model = WorkChainViewerModel()
+        ipw.dlink(
+            (self._model, "process_uuid"),
+            (model, "process_uuid"),
+        )
+        ipw.dlink(
+            (self._model, "monitor_counter"),
+            (model, "monitor_counter"),
+        )
+        self.node_view: WorkChainViewer = node_viewer(node, model=model)
+        self.node_view.render()
+        self.node_views[node.uuid] = self.node_view
+
     def _update_kill_button_layout(self):
         if not self.rendered:
             return

diff --git a/src/aiidalab_qe/app/result/structure/structure.py b/src/aiidalab_qe/app/result/structure/structure.py
@@ -5,7 +5,7 @@
 
 
 class StructureResults(ResultsPanel[StructureResultsModel]):
-    title = "Final geometry"
+    title = "Final Geometry"
     identifier = "structure"
 
     def render(self):

diff --git a/src/aiidalab_qe/app/result/viewer/viewer.py b/src/aiidalab_qe/app/result/viewer/viewer.py
@@ -3,7 +3,6 @@
 import ipywidgets as ipw
 import traitlets as tl
 
-from aiida import orm
 from aiidalab_qe.app.result.summary.model import WorkChainSummaryModel
 from aiidalab_qe.app.utils import get_entry_items
 from aiidalab_qe.common.panel import ResultsPanel
@@ -19,11 +18,7 @@
 class WorkChainViewer(ipw.VBox):
     _results_shown = tl.Set()
 
-    def __init__(self, node: orm.Node, model: WorkChainViewerModel, **kwargs):
-        if node.process_label != "QeAppWorkChain":
-            super().__init__()
-            return
-
+    def __init__(self, node, model: WorkChainViewerModel, **kwargs):
         from aiidalab_qe.common.widgets import LoadingWidget
 
         super().__init__(
@@ -32,6 +27,7 @@ def __init__(self, node: orm.Node, model: WorkChainViewerModel, **kwargs):
         )
 
         self._model = model
+        self._model.process_node = node
 
         self.rendered = False
 
@@ -49,9 +45,6 @@ def __init__(self, node: orm.Node, model: WorkChainViewerModel, **kwargs):
 
         self._fetch_plugin_results()
 
-        # HACK should be called from result step - fix!
-        self.render()
-
     def render(self):
         if self.rendered:
             return

diff --git a/tests/test_result.py b/tests/test_result.py
@@ -1,5 +1,3 @@
-import time
-
 from bs4 import BeautifulSoup
 
 from aiidalab_qe.app.main import App
@@ -33,8 +31,9 @@ def test_workchainview(generate_qeapp_workchain):
     workchain = generate_qeapp_workchain()
     workchain.node.seal()
     model = WorkChainViewerModel()
+    model.process_uuid = workchain.node.uuid
     viewer = WorkChainViewer(workchain.node, model=model)
-    time.sleep(3)
+    viewer.render()
     assert len(viewer.tabs.children) == 5
     assert viewer.tabs._titles["0"] == "Workflow Summary"  # type: ignore
     assert viewer.tabs._titles["1"] == "Final Geometry"  # type: ignore