Skip to content

Commit bda33c6

Browse files
fix mlff md test on ionic step data
1 parent 294d5ca commit bda33c6

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

tests/forcefields/test_md.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,8 @@ def test_ml_ff_md_maker(
115115
traj_file="md_traj.json.gz",
116116
traj_file_fmt="pmg",
117117
store_trajectory="partial",
118-
ionic_step_data=("energy", "forces", "stress", "mol_or_struct"),
118+
# check that `structure` alias to `mol_or_struct` works:
119+
ionic_step_data=("energy", "forces", "stress", "structure"),
119120
calculator_kwargs=calculator_kwargs,
120121
use_emmet_models=use_emmet_models,
121122
).make(structure)
@@ -135,10 +136,15 @@ def test_ml_ff_md_maker(
135136
# Check that the ionic steps have the expected physical properties
136137
assert all(
137138
key in step.model_dump()
138-
for key in ("energy", "forces", "stress", "mol_or_struct", "structure")
139+
for key in ("energy", "forces", "stress", "mol_or_struct")
139140
for step in task_doc.output.ionic_steps
140141
)
141142

143+
# `structure` aliases `mol_or_struct`
144+
assert all(
145+
step.structure == step.mol_or_struct for step in task_doc.output.ionic_steps
146+
)
147+
142148
# Check that the trajectory has expected physical properties
143149
assert task_doc.included_objects == ["trajectory"]
144150
assert len(task_doc.objects["trajectory"]) == n_steps + 1

0 commit comments

Comments
 (0)