From 29e00d3c9d6b9f75696091005216f127a3c0ac56 Mon Sep 17 00:00:00 2001
From: zhenyu wang <wzy@calypso.cn>
Date: Tue, 28 May 2024 08:45:03 +0800
Subject: [PATCH] fix: merge op output error

---
 dpgen2/superop/prep_run_calypso.py |  4 ++--
 tests/test_merge_caly_evo_step.py  | 11 +++++------
 tests/test_prep_run_caly.py        | 17 +++++++++--------
 3 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/dpgen2/superop/prep_run_calypso.py b/dpgen2/superop/prep_run_calypso.py
index da610ec2..3ba7334e 100644
--- a/dpgen2/superop/prep_run_calypso.py
+++ b/dpgen2/superop/prep_run_calypso.py
@@ -53,7 +53,7 @@ def __init__(
         self,
         name: str,
         prep_caly_input_op: Type[OP],
-        caly_evo_step_op: OPTemplate,
+        caly_evo_step_op: Type[OP],
         prep_caly_model_devi_op: Type[OP],
         run_caly_model_devi_op: Type[OP],
         prep_config: dict = normalize_step_dict({}),
@@ -143,7 +143,7 @@ def _prep_run_caly(
     prep_run_caly_steps: Steps,
     step_keys: Dict[str, Any],
     prep_caly_input_op: Type[OP],
-    caly_evo_step_op: OPTemplate,
+    caly_evo_step_op: Type[OP],
     prep_caly_model_devi_op: Type[OP],
     run_caly_model_devi_op: Type[OP],
     prep_config: dict = normalize_step_dict({}),
diff --git a/tests/test_merge_caly_evo_step.py b/tests/test_merge_caly_evo_step.py
index 4194f51b..00f18160 100644
--- a/tests/test_merge_caly_evo_step.py
+++ b/tests/test_merge_caly_evo_step.py
@@ -270,12 +270,11 @@ def setUp(self):
         )
 
     def tearDown(self):
-        pass
-        # shutil.rmtree(self.work_dir, ignore_errors=True)
-        # for i in Path().glob("caly-evo-step-*"):
-        #     shutil.rmtree(i, ignore_errors=True)
-        # for i in Path().glob("caly_task*"):
-        #     shutil.rmtree(i, ignore_errors=True)
+        shutil.rmtree(self.work_dir, ignore_errors=True)
+        for i in Path().glob("caly-evo-step-*"):
+            shutil.rmtree(i, ignore_errors=True)
+        for i in Path().glob("caly_task*"):
+            shutil.rmtree(i, ignore_errors=True)
 
     def test_caly_evo_step(self):
         steps = CalyEvoStepMerge(
diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py
index d555d418..547e844d 100644
--- a/tests/test_prep_run_caly.py
+++ b/tests/test_prep_run_caly.py
@@ -69,6 +69,9 @@
     BaseExplorationTaskGroup,
     ExplorationTask,
 )
+from dpgen2.op.caly_evo_step_merge import (
+    CalyEvoStepMerge,
+)
 from dpgen2.op.prep_caly_dp_optim import (
     PrepCalyDPOptim,
 )
@@ -81,9 +84,6 @@
 from dpgen2.op.run_caly_model_devi import (
     RunCalyModelDevi,
 )
-from dpgen2.superop.caly_evo_step import (
-    CalyEvoStep,
-)
 from dpgen2.superop.prep_run_calypso import (
     PrepRunCaly,
 )
@@ -149,11 +149,12 @@ def tearDown(self):
             shutil.rmtree(i, ignore_errors=True)
 
     def test(self):
-        caly_evo_step_op = CalyEvoStep(
-            "caly-evo-run",
-            MockedCollRunCaly,
-            PrepCalyDPOptim,
-            MockedRunCalyDPOptim,
+        caly_evo_step_op = CalyEvoStepMerge(
+            mode="debug",
+            name="caly-evo-step",
+            collect_run_caly=MockedCollRunCaly,
+            prep_dp_optim=PrepCalyDPOptim,
+            run_dp_optim=MockedRunCalyDPOptim,
             prep_config=prep_default_config,
             run_config=run_default_config,
             upload_python_packages=upload_python_packages,