SSH key prepare in advance and ssh set after

aiidalab_widgets_base/computational_resources.py

@@ -450,26 +450,30 @@ def _write_ssh_config(self, private_key_abs_fname=None):
                 file.write(f"  IdentityFile {private_key_abs_fname}\n")
             file.write("  ServerAliveInterval 5\n")
 
-    def _on_setup_ssh_button_pressed(self, _=None):
+    def key_pair_prepare(self):
+        """Prepare key pair for the ssh connection."""
         # Always start by generating a key pair if they are not present.
         self._ssh_keygen()
 
         # If hostname & username are not provided - do not do anything.
         if self.hostname.value == "":  # check hostname
-            self.message = "Please specify the computer hostname (for SSH)"
-            return False
+            message = "Please specify the computer name (for SSH)"
+
+            raise ValueError(message)
 
         if self.username.value == "":  # check username
-            self.message = "Please specify your SSH username."
-            return False
+            message = "Please specify your SSH username."
+
+            raise ValueError(message)
 
         private_key_abs_fname = None
         if self._verification_mode.value == "private_key":
             # unwrap private key file and setting temporary private_key content
             private_key_abs_fname, private_key_content = self._private_key
             if private_key_abs_fname is None:  # check private key file
-                self.message = "Please upload your private key file."
-                return False
+                message = "Please upload your private key file."
+
+                raise ValueError(message)
 
             # Write private key in ~/.ssh/ and use the name of upload file,
             # if exist, generate random string and append to filename then override current name.
@@ -478,10 +482,21 @@ def _on_setup_ssh_button_pressed(self, _=None):
         if not self._is_in_config():
             self._write_ssh_config(private_key_abs_fname=private_key_abs_fname)
 
-        # Copy public key on the remote computer.
+    def thread_ssh_copy_id(self):
+        """Copy public key on the remote computer, on a separate thread."""
         ssh_connection_thread = threading.Thread(target=self._ssh_copy_id)
         ssh_connection_thread.start()
 
+    def _on_setup_ssh_button_pressed(self, _=None):
+        """Setup ssh connection."""
+        try:
+            self.key_pair_prepare()
+        except ValueError as exc:
+            self.message = str(exc)
+            return
+
+        self.thread_ssh_copy_id()
+
     def _ssh_copy_id(self):
         """Run the ssh-copy-id command and follow it until it is completed."""
         timeout = 10
@@ -1654,9 +1669,16 @@ def _on_quick_setup(self, _=None):
             code_setup = self.template_variables_code.filled_templates
             self.code_setup = code_setup
 
+        # Setup the computer and code.
+        try:
+            self.ssh_computer_setup.key_pair_prepare()
+        except ValueError as exc:
+            self.message = f"Invalid inputs: {exc}"
+
         if self.aiida_computer_setup.on_setup_computer():
             self.aiida_code_setup.on_setup_code()
-        self.ssh_computer_setup._on_setup_ssh_button_pressed()
+
+        self.ssh_computer_setup.thread_ssh_copy_id
 
     def _on_setup_computer_success(self):
         """Callback that is called when the computer is successfully set up."""

tests/test_computational_resources.py

@@ -350,7 +350,6 @@ def test_template_variables_widget_multi_template_variables():
     assert w.filled_templates["default_calc_job_plugin"] == "quantumespresso.ph"
 
 
-@pytest.mark.usefixtures("aiida_profile_clean")
 def test_template_variables_widget_help_text_disappear_if_no_template_str():
     """This test when the template string is update to without template field, the help text should disappear."""
     w = computational_resources.TemplateVariablesWidget()
@@ -388,7 +387,7 @@ def test_quick_setup_widget():
     # Test message is update correctly. By click setup button without filling in any information.
     w._on_quick_setup()
 
-    assert "Please specify the computer hostname (for SSH)" in w.message
+    assert "Please specify the computer name" in w.message
 
     # 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.
@@ -429,15 +428,11 @@ def test_quick_setup_widget():
             sub_widget = mapping_variable.widget
             sub_widget.value = "pw"
 
-    w._on_quick_setup()
-
-    assert "Please specify your SSH username" in w.message
-
     w.ssh_computer_setup.username.value = "aiida"
 
     w._on_quick_setup()
 
-    assert "Generating SSH key pair" in w.message
+    assert "created" in w.message
     assert w.success
 
     # test select new resource reset the widget, success trait, and message trait.