diff --git a/dpgen2/op/prep_relax.py b/dpgen2/op/prep_relax.py index 1ee2869a..40a25312 100644 --- a/dpgen2/op/prep_relax.py +++ b/dpgen2/op/prep_relax.py @@ -41,7 +41,11 @@ def execute( ncifs = len(ip["cifs"]) config = ip["expl_config"] group_size = config["relax_group_size"] - ntasks = int(ncifs / group_size) + if group_size <= 0: + raise ValueError("relax_group_size must be greater than zero") + # Ceiling division keeps a final partial group instead of silently + # dropping CIFs when their count is not divisible by group_size. + ntasks = (ncifs + group_size - 1) // group_size task_paths = [] for i in range(ntasks): task_dir = Path("task.%06d" % i) diff --git a/tests/op/test_prep_relax.py b/tests/op/test_prep_relax.py index 83dcc2b7..b8bd60a1 100644 --- a/tests/op/test_prep_relax.py +++ b/tests/op/test_prep_relax.py @@ -15,32 +15,64 @@ class TestPrepRelax(unittest.TestCase): - def testPrepRelax(self): + def setUp(self): + self.cifs = [] + + def _make_cifs(self, count): cifs = [] - for i in range(4): + for i in range(count): p = Path("%i.cif" % i) p.write_text("Mocked cif.") cifs.append(p) + self.cifs.extend(cifs) + return cifs + + def _run_prep_relax(self, ncifs, group_size): op_in = OPIO( { "expl_config": { - "relax_group_size": 2, + "relax_group_size": group_size, }, - "cifs": cifs, + "cifs": self._make_cifs(ncifs), } ) op = PrepRelax() - op_out = op.execute(op_in) + return op.execute(op_in) + + def test_prep_relax(self): + op_out = self._run_prep_relax(4, 2) self.assertEqual(op_out["ntasks"], 2) self.assertEqual(len(op_out["task_paths"]), 2) for i, task_path in enumerate(op_out["task_paths"]): self.assertEqual(str(task_path), "task.%06d" % i) self.assertEqual(len(list(task_path.iterdir())), 2) + def test_keeps_partial_final_group(self): + """Assign every CIF when the final task is not a full group.""" + op_out = self._run_prep_relax(5, 2) + + self.assertEqual(op_out["ntasks"], 3) + self.assertEqual( + [len(list(task_path.iterdir())) for task_path in op_out["task_paths"]], + [2, 2, 1], + ) + + def test_creates_task_when_group_is_larger_than_input(self): + op_out = self._run_prep_relax(1, 2) + + self.assertEqual(op_out["ntasks"], 1) + self.assertEqual(len(list(op_out["task_paths"][0].iterdir())), 1) + + def test_rejects_non_positive_group_size(self): + for group_size in (0, -1): + with self.subTest(group_size=group_size): + with self.assertRaisesRegex(ValueError, "greater than zero"): + self._run_prep_relax(1, group_size) + def tearDown(self): - for i in range(2): - if os.path.isdir("task.%06d" % i): - shutil.rmtree("task.%06d" % i) - for i in range(4): - if os.path.isfile("%s.cif" % i): - os.remove("%s.cif" % i) + for task_path in Path().glob("task.[0-9][0-9][0-9][0-9][0-9][0-9]"): + if task_path.is_dir(): + shutil.rmtree(task_path) + for cif in self.cifs: + if cif.is_file(): + os.remove(cif)