From 25f7a690d006552f6ab166e68f6317c960e32aca Mon Sep 17 00:00:00 2001
From: Edan Bainglass <edan.bainglass@gmail.com>
Date: Fri, 10 Jan 2025 15:27:40 +0000
Subject: [PATCH] Minor changes

---
 .../app/result/components/status/tree.py      | 43 ++++++++-----------
 tests/test_status.py                          |  6 +--
 2 files changed, 20 insertions(+), 29 deletions(-)

diff --git a/src/aiidalab_qe/app/result/components/status/tree.py b/src/aiidalab_qe/app/result/components/status/tree.py
index 2091e499c..1dde3915b 100644
--- a/src/aiidalab_qe/app/result/components/status/tree.py
+++ b/src/aiidalab_qe/app/result/components/status/tree.py
@@ -95,34 +95,27 @@ def _collapse_all(self, _):
 
 
 class ProcessTreeNode(ipw.VBox, t.Generic[ProcessNodeType]):
-    _MAPPING = {
+    _TITLE_MAPPING = {
         "QeAppWorkChain": "Quantum ESPRESSO app workflow",
+        "BandsWorkChain": "Electronic band structure workflow",
         "PwBandsWorkChain": "Electronic band structure workflow",
         "ProjwfcBandsWorkChain": "Electronic band structure workflow",
+        "PwRelaxWorkChain": "Structure relaxation workflow",
         "PdosWorkChain": "Projected density of states workflow",
-        "DosCalculation": "Compute density of states",
-        "ProjwfcCalculation": "Compute projections",
-        "XpsWorkChain": "X-ray photoelectron spectroscopy workflow",
-        "XspectraCrystalWorkChain": "X-ray absorption spectroscopy workflow",
-    }
-
-    _PW_MAPPING = {
-        "scf": {
-            "PwBaseWorkChain": "SCF workflow",
-            "PwCalculation": "Run SCF cycle",
-        },
-        "nscf": {
-            "PwBaseWorkChain": "NSCF workflow",
-            "PwCalculation": "Run NSCF cycle",
+        "PwBaseWorkChain": {
+            "scf": "SCF workflow",
+            "nscf": "NSCF workflow",
+            "bands": "Bands workflow",
+            "relax": "Structure relaxation workflow",
         },
-        "bands": {
-            "PwBaseWorkChain": "Bands workflow",
-            "PwCalculation": "Compute bands",
-        },
-        "relax": {
-            "PwBaseWorkChain": "Structure relaxation workflow",
-            "PwCalculation": "Optimize structure geometry",
+        "PwCalculation": {
+            "scf": "Run SCF cycle",
+            "nscf": "Run NSCF cycle",
+            "bands": "Compute bands",
+            "relax": "Optimize structure geometry",
         },
+        "DosCalculation": "Compute density of states",
+        "ProjwfcCalculation": "Compute projections",
     }
 
     def __init__(
@@ -192,10 +185,8 @@ def _get_human_readable_title(self):
         if label in ("PwBaseWorkChain", "PwCalculation"):
             inputs = node.inputs.pw if label == "PwBaseWorkChain" else node.inputs
             calculation: str = inputs.parameters.get_dict()["CONTROL"]["calculation"]
-            mapping = self._PW_MAPPING[calculation]
-        else:
-            mapping = self._MAPPING
-        return mapping.get(label, label)
+            return self._TITLE_MAPPING.get(label, {}).get(calculation, label)
+        return self._TITLE_MAPPING.get(label, label)
 
 
 class ProcessTreeBranches(ipw.VBox):
diff --git a/tests/test_status.py b/tests/test_status.py
index 5abb9f0a1..64e0a4e53 100644
--- a/tests/test_status.py
+++ b/tests/test_status.py
@@ -135,7 +135,7 @@ def test_render(self):
     def test_update_on_process_change(self, mock_qeapp_workchain):
         self.model.process_uuid = mock_qeapp_workchain.uuid
         assert self.tree.rendered
-        human_label = ProcessTreeNode._MAPPING["QeAppWorkChain"]
+        human_label = ProcessTreeNode._TITLE_MAPPING["QeAppWorkChain"]
         assert self.tree.trunk.label.value == human_label
         assert not self.tree.trunk.collapsed
         assert len(self.tree.trunk.branches) == 1
@@ -146,7 +146,7 @@ def test_workchain_node(self):
         assert workchain_node.level == 1
         assert workchain_node.emoji.value == "✅"
         assert workchain_node.state.value == "finished"
-        human_label = ProcessTreeNode._PW_MAPPING["relax"]["PwBaseWorkChain"]
+        human_label = ProcessTreeNode._TITLE_MAPPING["PwBaseWorkChain"]["relax"]
         assert workchain_node.label.value == human_label
         assert isinstance(workchain_node.label, ipw.HTML)
         assert workchain_node.collapsed
@@ -186,7 +186,7 @@ def test_calculation_node(self):
         assert calculation_node.level == 2
         assert calculation_node.emoji.value == "✅"
         assert calculation_node.state.value == "finished"
-        human_label = ProcessTreeNode._PW_MAPPING["relax"]["PwCalculation"]
+        human_label = ProcessTreeNode._TITLE_MAPPING["PwCalculation"]["relax"]
         assert isinstance(calculation_node.label, ipw.Button)
         assert calculation_node.label.description == human_label