From fe650ec5ca47c0f307c6a128858d8ea2ada3ff55 Mon Sep 17 00:00:00 2001
From: ElliottKasoar <ElliottKasoar@users.noreply.github.com>
Date: Wed, 27 Nov 2024 18:09:56 +0000
Subject: [PATCH] Test abstract model

---
 tests/test_abstract_model.py | 55 +++++++++++++++++++++++++++---------
 1 file changed, 41 insertions(+), 14 deletions(-)

diff --git a/tests/test_abstract_model.py b/tests/test_abstract_model.py
index 1e91471..33840a2 100644
--- a/tests/test_abstract_model.py
+++ b/tests/test_abstract_model.py
@@ -11,24 +11,51 @@
 def extxyz_file():
     return StringIO(
         """2
-        Properties=species:S:1:pos:R:3:forces:R:3 energy=-100 pbc="F F F"
-        Si  0.0  0.0  0.0  0.46610144  0.6694725  -0.49123398
-        Si  0.0  0.0  0.0  -0.15967495  -0.51654063  -0.67342865
+        Properties=species:S:1:pos:R:3:forces:R:3 energy=-1 pbc="F T F"
+        Si  0.0  0.0  0.0  0.4  0.6  -0.4
+        Si  0.0  0.0  0.0  -0.1  -0.5  -0.6
         """
     )
 
 
 def test_from_atoms(extxyz_file):
+    """Test extracting data from ASE Atoms object."""
+    expected_forces = np.array([[0.4, 0.6, -0.4], [-0.1, -0.5, -0.6]])
+    expected_stress = np.array([-1.0, -1.0, -1.0, -2.1, 2.0, 1.8])
+
     atoms = read(extxyz_file, format="extxyz")
-    expected_forces = np.array(
-        [
-            [0.46610144,  0.6694725 , -0.49123398],
-            [-0.15967495, -0.51654063, -0.67342865],
-        ]
-    )
-    # atoms.calc.results["forces"] = forces
+    atoms.calc.results["stress"] = expected_stress
     data = AbstractModel.from_atoms(atoms)
-    assert data is not None
-    assert data["pbc"] == [False, False, False]
-    assert data["energy"] == -100
-    assert data["forces"] == expected_forces
+
+    # Test info
+    assert data["pbc"] == [False, True, False]
+    assert data["n_atoms"] == 2
+    assert len(data["cell"]) == 3
+    assert all(arr == [0.0, 0.0, 0.0] for arr in data["cell"])
+    assert data["formula"] == "Si2"
+
+    # Test results
+    assert data["energy"] == -1
+    assert data["forces"] == pytest.approx(expected_forces)
+    assert data["stress"] == pytest.approx(expected_stress)
+
+
+def test_from_atoms_no_calc(extxyz_file):
+    """Test extracting data from ASE Atoms object without results."""
+    expected_stress = np.array([-1.0, -1.0, -1.0, -2.1, 2.0, 1.8])
+
+    atoms = read(extxyz_file, format="extxyz")
+    atoms.calc.results["stress"] = expected_stress
+    data = AbstractModel.from_atoms(atoms, store_calc=False)
+
+    # Test info
+    assert data["pbc"] == [False, True, False]
+    assert data["n_atoms"] == 2
+    assert len(data["cell"]) == 3
+    assert all(arr == [0.0, 0.0, 0.0] for arr in data["cell"])
+    assert data["formula"] == "Si2"
+
+    # Test results
+    assert "energy" not in data
+    assert "forces" not in data
+    assert "stress" not in data