From c3e45aebeece8b1076ecc930b4af336134efd845 Mon Sep 17 00:00:00 2001
From: Yuan Chiang <cyrusyc@berkeley.edu>
Date: Mon, 7 Oct 2024 22:33:13 -0700
Subject: [PATCH] refactor test eos

---
 tests/test_eos.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/tests/test_eos.py b/tests/test_eos.py
index 8b5e069..7843f01 100644
--- a/tests/test_eos.py
+++ b/tests/test_eos.py
@@ -1,18 +1,20 @@
-import numpy as np
+import pytest
 from ase.build import bulk
 
+from mlip_arena.models.utils import MLIPEnum
 from mlip_arena.tasks.eos.run import fit as EOS
 
 atoms = bulk("Cu", "fcc", a=3.6)
 
-def test_eos():
+@pytest.mark.parametrize("model", [MLIPEnum["MACE-MP(M)"]])
+def test_eos(model: MLIPEnum):
     """
     Test EOS prefect workflow with a simple cubic lattice.
     """
 
     result = EOS(
         atoms=atoms,
-        calculator_name="MACE-MP(M)",
+        calculator_name=model.name,
         calculator_kwargs={},
         device=None,
         optimizer="BFGSLineSearch",
@@ -20,10 +22,10 @@ def test_eos():
         filter="FrechetCell",
         filter_kwargs=None,
         criterion=dict(
-            fmax=0.01,
+            fmax=0.1,
         ),
         max_abs_strain=0.1,
-        npoints=11,
+        npoints=6,
     )
 
     print(result)