fix generator

atomind-ai · Dec 22, 2024 · 5ec1a26 · 5ec1a26
1 parent a8f13fa
commit 5ec1a26
Showing 1 changed file with 14 additions and 13 deletions.
diff --git a/tests/test_mof.py b/tests/test_mof.py
@@ -9,17 +9,18 @@
 
 @pytest.mark.parametrize("model", [MLIPEnum["MACE-MP(M)"]])
 def test_widom_insertion(model: MLIPEnum):
-    atoms = get_atoms_from_db("mof.db")
     with prefect_test_harness():
-        result = widom_insertion.with_options(
-            refresh_cache=True,
-        )(
-            structure=atoms,
-            gas=molecule("CO2"),
-            calculator_name=model.name,
-            num_insertions=10,
-        )
-        assert isinstance(result, dict)
-        assert isinstance(result["henry_coefficient"][0], float)
-        assert isinstance(result["averaged_interaction_energy"][0], float)
-        assert isinstance(result["heat_of_adsorption"][0], float)
+        for atoms in get_atoms_from_db("mof.db"):
+            result = widom_insertion.with_options(
+                refresh_cache=True,
+            )(
+                structure=atoms,
+                gas=molecule("CO2"),
+                calculator_name=model.name,
+                num_insertions=10,
+            )
+            assert isinstance(result, dict)
+            assert isinstance(result["henry_coefficient"][0], float)
+            assert isinstance(result["averaged_interaction_energy"][0], float)
+            assert isinstance(result["heat_of_adsorption"][0], float)
+            break # only test one MOF