update UTs for the change in exploration

wanghan-iapcm · Feb 8, 2022 · de180f1 · de180f1
1 parent b2af611
commit de180f1
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 19 deletions.
diff --git a/tests/exploration/test_conf_filter.py b/tests/exploration/test_conf_filter.py
@@ -2,7 +2,7 @@
 import numpy as np
 import unittest, dpdata
 from exploration.context import dpgen2
-from dpgen2.exploration.conf_filter import ConfFilter, ConfFilters
+from dpgen2.exploration.selector import ConfFilter, ConfFilters
 from fake_data_set import fake_system
 from mock import patch
 

diff --git a/tests/exploration/test_conf_selector_frame.py b/tests/exploration/test_conf_selector_frame.py
@@ -3,8 +3,10 @@
 import numpy as np
 import unittest
 from pathlib import Path
-from dpgen2.exploration.conf_selector_frame import ConfSelectorLammpsFrames
-from dpgen2.exploration.trust_level import TrustLevel
+from dpgen2.exploration.selector import (
+    TrustLevel,
+    ConfSelectorLammpsFrames,
+)
 
 class TestConfSelectorLammpsFrames(unittest.TestCase):
     def setUp(self):

diff --git a/tests/exploration/test_exploration_group.py b/tests/exploration/test_exploration_group.py
@@ -9,8 +9,10 @@
 except ModuleNotFoundError:
     # case of upload everything to argo, no context needed
     pass
-from dpgen2.exploration.group import CPTGroup
-from dpgen2.exploration.stage import ExplorationStage
+from dpgen2.exploration.task import (
+    CPTGroup, 
+    ExplorationStage,
+)
 from dpgen2.constants import lmp_conf_name, lmp_input_name
 from unittest.mock import Mock, patch
 
@@ -91,7 +93,7 @@ class TestCPTGroup(unittest.TestCase):
     # def setUp(self):
     #     self.mock_random = Mock()
 
-    @patch('dpgen2.exploration.lmp.lmp_input.random')
+    @patch('dpgen2.exploration.task.lmp.lmp_input.random')
     def test_npt(self, mock_random):
         mock_random.randrange.return_value = 1110
         self.confs = ['foo', 'bar']
@@ -127,7 +129,7 @@ def test_npt(self, mock_random):
                 in_template_npt % (self.tt[j_idx], self.pp[k_idx]),
             )
 
-    @patch('dpgen2.exploration.lmp.lmp_input.random')
+    @patch('dpgen2.exploration.task.lmp.lmp_input.random')
     def test_nvt(self, mock_random):
         mock_random.randrange.return_value = 1110
         self.confs = ['foo', 'bar']
@@ -162,7 +164,7 @@ def test_nvt(self, mock_random):
             )
 
 
-    @patch('dpgen2.exploration.lmp.lmp_input.random')
+    @patch('dpgen2.exploration.task.lmp.lmp_input.random')
     def test_nvt_sample(self, mock_random):
         mock_random.randrange.return_value = 1110
         self.confs = ['foo', 'bar']
@@ -243,8 +245,8 @@ def test_nvt_sample(self, mock_random):
             )
 
 
-    @patch('dpgen2.exploration.group.random.shuffle')
-    @patch('dpgen2.exploration.lmp.lmp_input.random.randrange')
+    @patch('dpgen2.exploration.task.npt_group.random.shuffle')
+    @patch('dpgen2.exploration.task.lmp.lmp_input.random.randrange')
     def test_nvt_sample_random(self, mock_randrange, mock_shuffle):
         mock_randrange.return_value = 1110
         mock_shuffle.side_effect = swap_element
@@ -331,7 +333,7 @@ class TestCPTStage(unittest.TestCase):
     # def setUp(self):
     #     self.mock_random = Mock()
 
-    @patch('dpgen2.exploration.lmp.lmp_input.random')
+    @patch('dpgen2.exploration.task.lmp.lmp_input.random')
     def test(self, mock_random):
         mock_random.randrange.return_value = 1110
         self.numb_model = 3

diff --git a/tests/exploration/test_exploration_scheduler.py b/tests/exploration/test_exploration_scheduler.py
@@ -16,11 +16,9 @@
     ConstTrustLevelStageScheduler,
     ExplorationScheduler,
 )
-from dpgen2.exploration.stage import ExplorationStage
 from dpgen2.exploration.report import ExplorationReport
-from dpgen2.exploration.task import ExplorationTaskGroup
-from dpgen2.exploration.trust_level import TrustLevel
-from dpgen2.exploration.conf_selector import TrustLevelConfSelector
+from dpgen2.exploration.task import ExplorationTaskGroup, ExplorationStage
+from dpgen2.exploration.selector import TrustLevel, TrustLevelConfSelector
 from mocked_ops import (
     MockedExplorationReport,
     MockedExplorationTaskGroup,

diff --git a/tests/mocked_ops.py b/tests/mocked_ops.py
@@ -39,11 +39,9 @@
 from dpgen2.op.run_vasp import RunVasp
 from dpgen2.op.collect_data import CollectData
 from dpgen2.op.select_confs import SelectConfs
-from dpgen2.exploration.conf_selector import ConfSelector
-from dpgen2.exploration.trust_level import TrustLevel
-from dpgen2.exploration.task import ExplorationTask, ExplorationTaskGroup
+from dpgen2.exploration.selector import TrustLevel, ConfSelector
+from dpgen2.exploration.task import ExplorationTask, ExplorationTaskGroup, ExplorationStage
 from dpgen2.exploration.report import ExplorationReport
-from dpgen2.exploration.stage import ExplorationStage
 from dpgen2.exploration.scheduler import ConstTrustLevelStageScheduler
 
 mocked_template_script = { 'seed' : 1024, 'data': [] }