more advanced quick_setup

aiidalab_widgets_base/computational_resources.py

@@ -668,19 +668,19 @@ def _reset(self):
         self.proxy_command.value = ""
 
     @tl.observe("ssh_config")
-    def _observe_ssh_config(self, _=None):
+    def _observe_ssh_config(self, change):
         """Pre-filling the input fields."""
-        if not self.ssh_config:
-            self._reset()
+        self._reset()
 
-        if "hostname" in self.ssh_config:
-            self.hostname.value = self.ssh_config["hostname"]
-        if "port" in self.ssh_config:
-            self.port.value = int(self.ssh_config["port"])
-        if "proxy_jump" in self.ssh_config:
-            self.proxy_jump.value = self.ssh_config["proxy_jump"]
-        if "proxy_command" in self.ssh_config:
-            self.proxy_command.value = self.ssh_config["proxy_command"]
+        new_ssh_config = change["new"]
+        if "hostname" in new_ssh_config:
+            self.hostname.value = new_ssh_config["hostname"]
+        if "port" in new_ssh_config:
+            self.port.value = int(new_ssh_config["port"])
+        if "proxy_jump" in new_ssh_config:
+            self.proxy_jump.value = new_ssh_config["proxy_jump"]
+        if "proxy_command" in new_ssh_config:
+            self.proxy_command.value = new_ssh_config["proxy_command"]
 
 
 class AiidaComputerSetup(ipw.VBox):
@@ -1541,7 +1541,7 @@ def __init__(self, default_calc_job_plugin=None, **kwargs):
         self.comp_resources_database.observe(self._on_select_code, names="code_setup")
 
         self.ssh_computer_setup = SshComputerSetup()
-        self.ssh_computer_setup.observe(self._on_ssh_computer_setup, names="ssh_config")
+        # self.ssh_computer_setup.observe(self._on_ssh_computer_setup, names="ssh_config")
         ipw.dlink(
             (self.ssh_computer_setup, "message"),
             (self, "message"),
@@ -1606,10 +1606,10 @@ def __init__(self, default_calc_job_plugin=None, **kwargs):
             **kwargs,
         )
 
-    def _on_ssh_computer_setup(self, change=None):
-        """Callback when the ssh config is set."""
-        # Update the ssh config.
-        self.computer_configure = change["new"]
+    # def _on_ssh_computer_setup(self, change=None):
+    #    """Callback when the ssh config is set."""
+    #    # Update the ssh config.
+    #    self.computer_setup_and_configure["configure"] = change["new"]
 
     def _on_template_variables_computer_setup_filled(self, change):
         """Callback when the template variables of computer are filled."""
@@ -1663,10 +1663,11 @@ def _on_select_computer(self, change):
         self.computer_setup_and_configure = new_setup_and_configure
 
         # ssh config need to sync hostname etc with resource database.
-        ssh_config = self._parse_ssh_config_from_computer_configure(
-            new_setup_and_configure["configure"]
+        self.ssh_computer_setup.ssh_config = (
+            self._parse_ssh_config_from_computer_configure(
+                new_setup_and_configure["configure"]
+            )
         )
-        self.ssh_computer_setup.ssh_config = ssh_config
 
         # decide whether to show the ssh password box widget.
         # Since for 2FA ssh credential, the password are not needed but set from
@@ -1797,9 +1798,13 @@ def reset(self):
         # reset sub widgets
         self.aiida_code_setup._reset()
         self.aiida_computer_setup._reset()
+
         self.ssh_computer_setup._reset()
         self.ssh_auth = None
 
+        # essential, since if not, the same computer_configure won't trigger the `_observe_ssh_config` callback.
+        self.ssh_computer_setup.ssh_config = {}
+
         # reset traits
         self.computer_setup_and_configure = {}
         self.code_setup = {}

tests/test_computational_resources.py

@@ -507,13 +507,14 @@ def test_quick_setup_widget():
 
     w.ssh_computer_setup.username.value = "aiida"
 
-    # import ipdb; ipdb.set_trace()
+    # XXX: since cscs is 2FA, test the password box is not displayed.
 
     w._on_quick_setup()
 
     assert "created" in w.message
     assert "pw" in w.message
     assert w.success
+    assert orm.load_code("pw-7.2@daint-mc")
 
     # test select new resource reset the widget, success trait, and message trait.
     w.comp_resources_database.reset()
@@ -531,13 +532,76 @@ def test_quick_setup_widget():
     assert w.template_variables_computer_setup._template_variables != {}
 
 
+@pytest.mark.usefixtures("aiida_profile_clean")
 def test_quick_setup_widget_for_password_configure(monkeypatch, tmp_path):
     """Test for computer configure with password as ssh auth.
     The ssh auth is password, thus will generate ssh key pair and try to upload the key
     """
+    from aiidalab_widgets_base.computational_resources import QuickSetupWidget
+
     # monkeypatch home so the ssh key is generated in the temporary directory
     monkeypatch.setenv("HOME", str(tmp_path))
 
+    w = QuickSetupWidget()
+
+    # Test select a new resource setup will update the output interface (e.g. ssh_config, computer_setup, code_setup)
+    # and the computer/code setup widget will be updated accordingly.
+    w.comp_resources_database.domain_selector.value = "merlin.psi.ch"
+    w.comp_resources_database.computer_selector.value = "cpu"
+    w.comp_resources_database.code_selector.value = "QE-7.0-exe-template"
+
+    # Fill in the computer name and trigger the setup button again, the message should be updated.
+    for (
+        key,
+        mapping_variable,
+    ) in w.template_variables_computer_setup._template_variables.items():
+        if key == "label":
+            sub_widget = mapping_variable.widget
+
+            # Test the default value is filled in correctly.
+            assert sub_widget.value == "merlin-cpu"
+
+    # XXX test the password box is displayed.
+
+    # Fill the computer configure template variables
+    for (
+        key,
+        mapping_variable,
+    ) in w.template_variables_computer_configure._template_variables.items():
+        if key == "username":
+            sub_widget = mapping_variable.widget
+            sub_widget.value = "aiida"
+
+    # Fill the code name
+    for key, mapping_variable in w.template_variables_code._template_variables.items():
+        if key == "code_binary_name":
+            sub_widget = mapping_variable.widget
+            sub_widget.value = "ph"
+
+            # select the other code and check the filled template is updated
+            sub_widget.value = "pw"
+
+    w.ssh_computer_setup.username.value = "aiida"
+
+    # The quick_setup with password auth will try connect which will timeout.
+    # Thus, mock the connect method to avoid the timeout.
+    monkeypatch.setattr(
+        "aiidalab_widgets_base.computational_resources.SshComputerSetup.thread_ssh_copy_id",
+        lambda _: None,
+    )
+    w._on_quick_setup()
+
+    assert w.success
+    # check the code is really created
+    assert orm.load_code("pw-7.0@merlin-cpu")
+
+    # The key pair will be generated to the temporary directory
+    # Check the content of the config is correct
+    with open(tmp_path / ".ssh" / "config") as f:
+        content = f.read()
+        assert "User aiida" in content
+        assert "Host merlin-l-01.psi.ch" in content
+
 
 @pytest.mark.usefixtures("aiida_profile_clean")
 def test_quick_setup_widget_computer_change_code_reset():