fix: fix calypso unittest

dpgen2/op/collect_run_caly.py

@@ -67,7 +67,7 @@ def get_input_sign(cls):
                     type=Path, optional=True
                 ),  # dir named results for evo
                 "opt_results_dir": Artifact(
-                    type=Path, optional=True
+                    type=List[Path], optional=True
                 ),  # dir contains POSCAR* CONTCAR* OUTCAR*
                 "qhull_input": Artifact(type=Path, optional=True),  # for vsc
             }
@@ -141,11 +141,16 @@ def execute(
         results = (
             ip["results"].resolve() if ip["results"] is not None else ip["results"]
         )
-        opt_results_dir = (
-            ip["opt_results_dir"].resolve()
-            if ip["opt_results_dir"] is not None
-            else ip["opt_results_dir"]
-        )
+        # opt_results_dir = (
+        #     ip["opt_results_dir"].resolve()
+        #     if ip["opt_results_dir"] is not None
+        #     else ip["opt_results_dir"]
+        # )
+        opt_results_dir = []
+        if ip["opt_results_dir"] is not None:
+            for temp in ip["opt_results_dir"]:
+                opt_results_dir.append(Path(temp).resolve())
+
         qhull_input = (
             ip["qhull_input"].resolve()
             if ip["qhull_input"] is not None

tests/op/test_collect_run_caly.py

@@ -102,7 +102,7 @@ def side_effect(*args, **kwargs):
                     "input_file": self.input_file,
                     "step": self.step_file,
                     "results": self.results_dir,
-                    "opt_results_dir": self.opt_results_dir,
+                    "opt_results_dir": [self.opt_results_dir],
                 }
             )
         )
@@ -138,7 +138,7 @@ def side_effect(*args, **kwargs):
                     "input_file": self.input_file,
                     "step": self.step_file,
                     "results": self.results_dir,
-                    "opt_results_dir": self.opt_results_dir,
+                    "opt_results_dir": [self.opt_results_dir],
                 }
             )
         )
@@ -169,7 +169,7 @@ def side_effect(*args, **kwargs):
                     "input_file": self.input_file,
                     "step": self.step_file,
                     "results": self.results_dir,
-                    "opt_results_dir": self.opt_results_dir,
+                    "opt_results_dir": [self.opt_results_dir],
                 }
             )
         )
@@ -203,7 +203,7 @@ def side_effect(*args, **kwargs):
                     "input_file": self.input_file,
                     "step": self.step_file,
                     "results": self.results_dir,
-                    "opt_results_dir": self.opt_results_dir,
+                    "opt_results_dir": [self.opt_results_dir],
                 }
             ),
         )

tests/op/test_run_caly_model_devi.py

@@ -143,7 +143,7 @@ def side_effect_1(*args, **kwargs):
         mocked_run_1.side_effect = side_effect_1
 
         def side_effect_2(*args, **kwargs):
-            return [[1, 1, 1, 1, 1, 1, 1]]
+            return [[1, 1, 1, 1, 1, 1, 1, 1]]
 
         mocked_run_2.side_effect = side_effect_2
 

tests/op/test_run_dp_optim.py

@@ -102,13 +102,6 @@ def side_effect(*args, **kwargs):
             Path(self.task_name) / "traj_results" / "0.3.traj" in list_traj_results_dir
         )
 
-        self.assertEqual(
-            Path(self.task_name) / calypso_run_opt_file, out["caly_run_opt_file"]
-        )
-        self.assertEqual(
-            Path(self.task_name) / calypso_check_opt_file, out["caly_check_opt_file"]
-        )
-
     @patch("dpgen2.op.run_dp_optim.run_command")
     def test_01_error(self, mocked_run):
         def side_effect(*args, **kwargs):

tests/test_caly_evo_step.py

@@ -51,6 +51,20 @@
 except ModuleNotFoundError:
     # case of upload everything to argo, no context needed
     pass
+from context import (
+    default_host,
+    default_image,
+    skip_ut_with_dflow,
+    skip_ut_with_dflow_reason,
+    upload_python_packages,
+)
+from mocked_ops import (
+    MockedCollRunCaly,
+    MockedPrepDPOptim,
+    MockedRunDPOptim,
+    mocked_numb_models,
+)
+
 from dpgen2.constants import (
     lmp_conf_name,
     lmp_input_name,
@@ -82,20 +96,6 @@
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
 
-from .context import (
-    default_host,
-    default_image,
-    skip_ut_with_dflow,
-    skip_ut_with_dflow_reason,
-    upload_python_packages,
-)
-from .mocked_ops import (
-    MockedCollRunCaly,
-    MockedPrepDPOptim,
-    MockedRunDPOptim,
-    mocked_numb_models,
-)
-
 default_config = normalize_step_dict(
     {
         "template_config": {
@@ -109,7 +109,6 @@
 )
 
 
-@unittest.skip("temp")
 class TestMockedCollRunCaly(unittest.TestCase):
     def setUp(self) -> None:
         self.config = {}
@@ -160,7 +159,6 @@ def test_mocked_coll_run_caly_00(self):
         self.assertTrue(out["results"] == Path(self.task_name).joinpath("results"))
 
 
-@unittest.skip("temp")
 class TestMockedRunDPOptim(unittest.TestCase):
     def setUp(self) -> None:
         self.config = {}
@@ -204,7 +202,7 @@ def test_mocked_run_dp_optim(self):
         )
 
         self.assertTrue(optim_results_dir, Path(self.task_name) / "optim_results_dir")
-        self.assertEqual(counts_optim_results_dir, 10)
+        self.assertEqual(counts_optim_results_dir, 15)
         self.assertEqual(counts_outcar_in_optim_results_dir, 5)
         self.assertTrue(
             Path(self.task_name) / "optim_results_dir" / "CONTCAR_4"
@@ -220,13 +218,6 @@ def test_mocked_run_dp_optim(self):
             Path(self.task_name) / "traj_results" / "3.traj", list_traj_results_dir
         )
 
-        self.assertEqual(
-            Path(self.task_name) / calypso_run_opt_file, out["caly_run_opt_file"]
-        )
-        self.assertEqual(
-            Path(self.task_name) / calypso_check_opt_file, out["caly_check_opt_file"]
-        )
-
 
 # @unittest.skip("temporary pass")
 @unittest.skipIf(skip_ut_with_dflow, skip_ut_with_dflow_reason)
@@ -281,6 +272,7 @@ def tearDown(self):
             shutil.rmtree(i, ignore_errors=True)
         for i in Path().glob("caly_task*"):
             shutil.rmtree(i, ignore_errors=True)
+        shutil.rmtree("upload", ignore_errors=True)
 
     @unittest.skip("only need to run test_01")
     def test_00(self):

tests/test_prep_run_caly.py

@@ -60,8 +60,8 @@
 )
 from mocked_ops import (
     MockedCollRunCaly,
-    MockedPrepRunDPOptim,
     MockedRunCalyModelDevi,
+    MockedRunDPOptim,
     mocked_numb_models,
 )
 
@@ -72,6 +72,9 @@
 from dpgen2.op.prep_caly_input import (
     PrepCalyInput,
 )
+from dpgen2.op.prep_dp_optim import (
+    PrepDPOptim,
+)
 from dpgen2.op.run_caly_model_devi import (
     RunCalyModelDevi,
 )
@@ -83,11 +86,22 @@
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
 
-default_config = normalize_step_dict(
+prep_default_config = normalize_step_dict(
+    {
+        "template_config": {
+            "image": default_image,
+        },
+    }
+)
+run_default_config = normalize_step_dict(
     {
         "template_config": {
             "image": default_image,
-        }
+        },
+        "template_slice_config": {
+            "group_size": 2,
+            "pool_size": 1,
+        },
     }
 )
 
@@ -105,7 +119,7 @@ def make_task_group_list(njobs):
 
 # @unittest.skip("temporary pass")
 @unittest.skipIf(skip_ut_with_dflow, skip_ut_with_dflow_reason)
-class TestCalyEvoStep(unittest.TestCase):
+class TestPrepRunCaly(unittest.TestCase):
     def setUp(self):
         self.expl_config = {}
         self.work_dir = Path("storge_files")
@@ -114,7 +128,9 @@ def setUp(self):
         self.nmodels = mocked_numb_models
         self.model_list = []
         for ii in range(self.nmodels):
-            model = self.work_dir.joinpath(f"model.{ii}.pb")
+            model_path = self.work_dir.joinpath(f"task.{ii}")
+            model_path.mkdir(parents=True, exist_ok=True)
+            model = model_path.joinpath(f"frozen_model.pb")
             model.write_text(f"model {ii}")
             self.model_list.append(model)
         self.models = upload_artifact(self.model_list)
@@ -127,24 +143,25 @@ def tearDown(self):
         shutil.rmtree(self.work_dir, ignore_errors=True)
         for i in Path().glob("prep-run-caly-step*"):
             shutil.rmtree(i, ignore_errors=True)
-        # shutil.rmtree("upload", ignore_errors=True)
+        shutil.rmtree("upload", ignore_errors=True)
 
     def test(self):
         caly_evo_step_op = CalyEvoStep(
             "caly-evo-run",
             MockedCollRunCaly,
-            MockedPrepRunDPOptim,
-            prep_config=default_config,
-            run_config=default_config,
+            PrepDPOptim,
+            MockedRunDPOptim,
+            prep_config=prep_default_config,
+            run_config=run_default_config,
             upload_python_packages=upload_python_packages,
         )
         prep_run_caly_op = PrepRunCaly(
             "prep-run-calypso",
             PrepCalyInput,
             caly_evo_step_op,
             MockedRunCalyModelDevi,
-            prep_config=default_config,
-            run_config=default_config,
+            prep_config=prep_default_config,
+            run_config=run_default_config,
             upload_python_packages=upload_python_packages,
         )
         prep_run_caly_step = Step(