From d3dc01d5ee57c9a9ba2b2429c7741cdd21187cff Mon Sep 17 00:00:00 2001
From: Zhenyu Wang <wzy@calypso.cn>
Date: Fri, 29 Mar 2024 09:43:44 +0800
Subject: [PATCH] fix: fix ut

---
 tests/test_prep_run_caly.py     | 13 +++++++------
 tests/test_prep_run_dp_train.py | 21 +++++++++++----------
 tests/test_prep_run_lmp.py      |  9 +++++----
 tests/test_prep_run_vasp.py     | 11 ++++++-----
 tests/test_select_confs.py      | 10 +++++-----
 5 files changed, 34 insertions(+), 30 deletions(-)

diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py
index 5daa600f..f1e315ee 100644
--- a/tests/test_prep_run_caly.py
+++ b/tests/test_prep_run_caly.py
@@ -58,12 +58,6 @@
     skip_ut_with_dflow_reason,
     upload_python_packages,
 )
-from mocked_ops import (
-    MockedCollRunCaly,
-    MockedRunCalyModelDevi,
-    MockedRunDPOptim,
-    mocked_numb_models,
-)
 
 from dpgen2.exploration.task import (
     BaseExplorationTaskGroup,
@@ -86,6 +80,13 @@
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
 
+from .mocked_ops import (
+    MockedCollRunCaly,
+    MockedRunCalyModelDevi,
+    MockedRunDPOptim,
+    mocked_numb_models,
+)
+
 prep_default_config = normalize_step_dict(
     {
         "template_config": {
diff --git a/tests/test_prep_run_dp_train.py b/tests/test_prep_run_dp_train.py
index 22eb40ba..87ebe541 100644
--- a/tests/test_prep_run_dp_train.py
+++ b/tests/test_prep_run_dp_train.py
@@ -49,16 +49,6 @@
     skip_ut_with_dflow_reason,
     upload_python_packages,
 )
-from mocked_ops import (
-    MockedModifyTrainScript,
-    MockedPrepDPTrain,
-    MockedRunDPTrain,
-    MockedRunDPTrainNoneInitModel,
-    make_mocked_init_data,
-    make_mocked_init_models,
-    mocked_numb_models,
-    mocked_template_script,
-)
 
 from dpgen2.constants import (
     train_task_pattern,
@@ -69,6 +59,17 @@
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
 
+from .mocked_ops import (
+    MockedModifyTrainScript,
+    MockedPrepDPTrain,
+    MockedRunDPTrain,
+    MockedRunDPTrainNoneInitModel,
+    make_mocked_init_data,
+    make_mocked_init_models,
+    mocked_numb_models,
+    mocked_template_script,
+)
+
 default_config = normalize_step_dict(
     {
         "template_config": {
diff --git a/tests/test_prep_run_lmp.py b/tests/test_prep_run_lmp.py
index 3b350240..20ca570f 100644
--- a/tests/test_prep_run_lmp.py
+++ b/tests/test_prep_run_lmp.py
@@ -51,10 +51,6 @@
     skip_ut_with_dflow_reason,
     upload_python_packages,
 )
-from mocked_ops import (
-    MockedRunLmp,
-    mocked_numb_models,
-)
 
 from dpgen2.constants import (
     lmp_conf_name,
@@ -80,6 +76,11 @@
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
 
+from .mocked_ops import (
+    MockedRunLmp,
+    mocked_numb_models,
+)
+
 default_config = normalize_step_dict(
     {
         "template_config": {
diff --git a/tests/test_prep_run_vasp.py b/tests/test_prep_run_vasp.py
index 77030bff..ac306375 100644
--- a/tests/test_prep_run_vasp.py
+++ b/tests/test_prep_run_vasp.py
@@ -50,11 +50,6 @@
     skip_ut_with_dflow_reason,
     upload_python_packages,
 )
-from mocked_ops import (
-    MockedPrepVasp,
-    MockedRunVasp,
-    mocked_incar_template,
-)
 
 from dpgen2.constants import (
     fp_task_pattern,
@@ -70,6 +65,12 @@
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
 
+from .mocked_ops import (
+    MockedPrepVasp,
+    MockedRunVasp,
+    mocked_incar_template,
+)
+
 default_config = normalize_step_dict(
     {
         "template_config": {
diff --git a/tests/test_select_confs.py b/tests/test_select_confs.py
index c8840eaf..576ef36f 100644
--- a/tests/test_select_confs.py
+++ b/tests/test_select_confs.py
@@ -45,16 +45,16 @@
 except ModuleNotFoundError:
     # case of upload everything to argo, no context needed
     pass
-from mocked_ops import (
+from dpgen2.op.select_confs import (
+    SelectConfs,
+)
+
+from .mocked_ops import (
     MockedConfSelector,
     MockedExplorationReport,
     MockedSelectConfs,
 )
 
-from dpgen2.op.select_confs import (
-    SelectConfs,
-)
-
 
 class TestMockedSelectConfs(unittest.TestCase):
     def setUp(self):