From 9f4e503fe6c5feefaf55793474547fbc3662bf32 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Tue, 13 Aug 2024 23:12:01 +0100
Subject: [PATCH] Fix imports

---
 src/aiidalab_qe/__main__.py                   | 16 +++++++++-------
 src/aiidalab_qe/app/configuration/advanced.py |  2 +-
 src/aiidalab_qe/app/configuration/pseudos.py  |  2 +-
 tests/conftest.py                             |  2 +-
 tests/test_cli.py                             |  4 ++--
 tests/test_configure.py                       |  2 +-
 tests/test_pseudo.py                          |  2 +-
 7 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/src/aiidalab_qe/__main__.py b/src/aiidalab_qe/__main__.py
index a0911719a..1a566d5b5 100644
--- a/src/aiidalab_qe/__main__.py
+++ b/src/aiidalab_qe/__main__.py
@@ -5,10 +5,6 @@
 
 import click
 
-from aiida import load_profile
-from aiidalab_qe.setup.setup_codes import codes_are_setup
-from aiidalab_qe.setup.setup_codes import install as install_qe_codes
-
 # The default profile name of AiiDAlab container.
 _DEFAULT_PROFILE = "default"
 
@@ -22,9 +18,12 @@ def cli():
 @click.option("-f", "--force", is_flag=True)
 @click.option("-p", "--profile", default=_DEFAULT_PROFILE)
 def install_qe(force, profile):
+    from aiida import load_profile
+    from aiidalab_qe.setup.setup_codes import codes_are_setup, install
+
     load_profile(profile)
     try:
-        for msg in install_qe_codes(force=force):
+        for msg in install(force=force):
             click.echo(msg)
         assert codes_are_setup()
         click.secho("Codes are setup!", fg="green")
@@ -45,7 +44,8 @@ def install_pseudos(profile, source):
     """Install pseudopotentials from a local folder if source is specified,
     otherwise download from remote repositories.
     """
-    from aiidalab_qe.common.setup_pseudos import install
+    from aiida import load_profile
+    from aiidalab_qe.setup.setup_pseudos import install
 
     load_profile(profile)
 
@@ -90,9 +90,11 @@ def download_pseudos(dest):
 )
 @click.option("-p", "--profile", default=_DEFAULT_PROFILE)
 def test_plugin(plugin_name, profile):
-    load_profile(profile)
+    from aiida import load_profile
     from aiidalab_qe.app.utils import test_plugin_functionality
 
+    load_profile(profile)
+
     try:
         success, message = test_plugin_functionality(plugin_name)
         if success:
diff --git a/src/aiidalab_qe/app/configuration/advanced.py b/src/aiidalab_qe/app/configuration/advanced.py
index 65a416960..db6c09b64 100644
--- a/src/aiidalab_qe/app/configuration/advanced.py
+++ b/src/aiidalab_qe/app/configuration/advanced.py
@@ -18,8 +18,8 @@
 from aiida_quantumespresso.workflows.pw.base import PwBaseWorkChain
 from aiidalab_qe.app.parameters import DEFAULT_PARAMETERS
 from aiidalab_qe.common.panel import Panel
-from aiidalab_qe.common.setup_pseudos import PseudoFamily
 from aiidalab_qe.common.widgets import HubbardWidget
+from aiidalab_qe.setup.setup_pseudos import PseudoFamily
 
 from .pseudos import PseudoFamilySelector, PseudoSetter
 
diff --git a/src/aiidalab_qe/app/configuration/pseudos.py b/src/aiidalab_qe/app/configuration/pseudos.py
index 2a46a9bce..96ceffc26 100644
--- a/src/aiidalab_qe/app/configuration/pseudos.py
+++ b/src/aiidalab_qe/app/configuration/pseudos.py
@@ -11,7 +11,7 @@
 from aiida.plugins import DataFactory, GroupFactory
 from aiida_quantumespresso.workflows.pw.base import PwBaseWorkChain
 from aiidalab_qe.app.parameters import DEFAULT_PARAMETERS
-from aiidalab_qe.common.setup_pseudos import (
+from aiidalab_qe.setup.setup_pseudos import (
     PSEUDODOJO_VERSION,
     SSSP_VERSION,
     PseudoFamily,
diff --git a/tests/conftest.py b/tests/conftest.py
index b7c83a187..148f2a29b 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -5,7 +5,7 @@
 import tempfile
 
 import pytest
-from aiidalab_qe.common.setup_pseudos import SSSP_VERSION
+from aiidalab_qe.setup.setup_pseudos import SSSP_VERSION
 
 from aiida import orm
 
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 91bb0668f..f3b284938 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -24,7 +24,7 @@ def test_download_and_install_pseudos(tmp_path, aiida_profile, monkeypatch):
 
     Note: this test is slow, it takes about ~30 seconds to run.
     """
-    from aiidalab_qe.common.setup_pseudos import (
+    from aiidalab_qe.setup.setup_pseudos import (
         PSEUDODOJO_VERSION,
         SSSP_VERSION,
         pseudos_to_install,
@@ -39,7 +39,7 @@ def test_download_and_install_pseudos(tmp_path, aiida_profile, monkeypatch):
     }
     # mock the EXPECTED_PSEUDOS
     monkeypatch.setattr(
-        "aiidalab_qe.common.setup_pseudos.EXPECTED_PSEUDOS",
+        "aiidalab_qe.setup.setup_pseudos.EXPECTED_PSEUDOS",
         MOCK_EXPECTED_PSEUDOS,
     )
 
diff --git a/tests/test_configure.py b/tests/test_configure.py
index 0af35c976..61b339577 100644
--- a/tests/test_configure.py
+++ b/tests/test_configure.py
@@ -1,4 +1,4 @@
-from aiidalab_qe.common.setup_pseudos import PSEUDODOJO_VERSION, SSSP_VERSION
+from aiidalab_qe.setup.setup_pseudos import PSEUDODOJO_VERSION, SSSP_VERSION
 
 
 def test_protocol():
diff --git a/tests/test_pseudo.py b/tests/test_pseudo.py
index a85c34e86..70eba0dd2 100644
--- a/tests/test_pseudo.py
+++ b/tests/test_pseudo.py
@@ -1,5 +1,5 @@
 import pytest
-from aiidalab_qe.common.setup_pseudos import (
+from aiidalab_qe.setup.setup_pseudos import (
     PSEUDODOJO_VERSION,
     SSSP_VERSION,
     _construct_cmd,