From cf69e2d193fb23c8bed5ad0e8a3f22a0a0e634f6 Mon Sep 17 00:00:00 2001 From: Jusong Yu Date: Thu, 5 Oct 2023 09:47:05 +0000 Subject: [PATCH] more advanced quick_setup --- .../computational_resources.py | 43 ++++++------ tests/test_computational_resources.py | 67 ++++++++++++++++++- 2 files changed, 89 insertions(+), 21 deletions(-) diff --git a/aiidalab_widgets_base/computational_resources.py b/aiidalab_widgets_base/computational_resources.py index 7efd68941..e4ca2c763 100644 --- a/aiidalab_widgets_base/computational_resources.py +++ b/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 = {} diff --git a/tests/test_computational_resources.py b/tests/test_computational_resources.py index 0be307191..4c5b481ba 100644 --- a/tests/test_computational_resources.py +++ b/tests/test_computational_resources.py @@ -30,7 +30,6 @@ def test_ssh_computer_setup_widget(monkeypatch, tmp_path): widget.username.value = "aiida" # Write the information to ~/.ssh/config and check that it is there. - # XXX make this test independent of the user's ~/.ssh/config assert widget._is_in_config() is False widget._write_ssh_config() assert widget._is_in_config() is True @@ -507,13 +506,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 +531,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():